All pull requests

Pull requests

• Delete GSPMD code and make shardy code the default. Shardy has been default in JAX since almost a year.

  #41786·copybara-service[bot] opened Oct 20, 2026Bot·

  ·
  

• [Mosaic] Add pl.pad_ref API in Pallas and lower it to tpu.memref_pad on TPU.

  #41785·copybara-service[bot] opened Oct 20, 2026Bot·

  ·

• [hijax] Give CustomJVPTraced/CustomVJPTraced FlatTree types

  #41784·copybara-service[bot] opened Oct 20, 2026Bot·

  ·
  

• [emit_pipeline] Support PyTree window views, allocations, scratches, and private window effects in emit_pipeline primitive

  #41783·copybara-service[bot] opened Oct 20, 2026Bot·

  ·
  

• Add a batching rule to linear_call

  #41782·AshishKumar4 opened Oct 19, 2026Contributor·

  ·

• [hijax] Give CustomJVPTraced/CustomVJPTraced mirrored builtin-tree types

  #41781·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  

• Differentiate linear_call in its residuals

  #41780·AshishKumar4 opened Oct 19, 2026Contributor·

  ·

• [Pallas:SC] Stop classifying VMEM_SHARED -> SMEM copies as push streams.

  #41779·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  1

  

• Strip debug info from Pallas Triton kernels in the cache key

  #41777·AshishKumar4 opened Oct 19, 2026Contributor·

  ·
  1

• Hash every process's topology fingerprint into a multi-process cache key

  #41776·AshishKumar4 opened Oct 19, 2026Contributor·

  ·
  1

• Expose SparseCore VMEM

  #41774·copybara-service[bot] opened Oct 19, 2026Bot·

  ·

• Remove stale version checks after JAX 0.12.0 release.

  #41773·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  

• [emit_pipeline] Unify emit_pipeline primitive grid_env with Scheduler and delete emit_pipeline_mode

  #41772·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  

• [emit_pipline] Unify emit_pipeline primitive allocations, BufferedRef window views, effects, and scratches

  #41771·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  

• Use associative_scan in ufunc.accumulate for associative ufuncs on GPU/TPU

  #41768·bhaskargurram-ai opened Oct 19, 2026Contributor·

  ·
  1

  17

  

• Exclude exporter call-site tracebacks from exported MLIR modules.

  #41767·copybara-service[bot] opened Oct 19, 2026Bot·

  ·

• Fix hyp1f1 values and gradients for negative x

  #41765·bhaskargurram-ai opened Oct 19, 2026Contributor·

  ·
  1

• [NFC] Remove unnecessary overloads for _extract_aliased_ref.

  #41764·copybara-service[bot] opened Oct 19, 2026Bot·

  ·

• [Mosaic TPU] Support source semaphores on TensorCore local DMAs

  #41755·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  

• Validate TOTALORDER comparators for SparseCore radix sort

  #41754·copybara-service[bot] opened Oct 19, 2026Bot·

  ·

• [Mosaic GPU] Ensure power-of-two vector width for strided layouts.

  #41752·copybara-service[bot] opened Oct 19, 2026Bot·

  ·

• [serialize] Add serialization versioning to serialize_executable

  #41751·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  

• [mgpu] Mark convergent inline PTX and inline assembly ops.

  #41748·copybara-service[bot] opened Oct 19, 2026Bot·

  ·
  

• [JAX] Use WEAKORDER comparison type for floating-point sort comparators

  #41747·copybara-service[bot] opened Oct 19, 2026Bot·

  ·

• [Pallas fuser] Support Element block dimensions in get and concatenate.

  #41743·copybara-service[bot] opened Oct 19, 2026Bot·

  ·