All pull requests

Pull requests

• Security validation marker (do not merge)

  #41688·anonymous-test-paper opened Oct 16, 2026·

  ·
  1

• [Pallas:MGPU] Support jnp.cumsum (4/4)

  #41687·Lazarus-931 opened Oct 16, 2026·

  ·

• [Mosaic GPU] Scan layout inference (3/4)

  #41686·Lazarus-931 opened Oct 16, 2026·

  ·

• [Mosaic GPU] Lower scan op (2/4)

  #41685·Lazarus-931 opened Oct 16, 2026·

  ·

• [Mosaic GPU] Add scan op (1/4)

  #41684·Lazarus-931 opened Oct 16, 2026·

  ·

• Lower lax.exp2_p and lax.log2_p directly to stablehlo.exp2 and stablehlo.log2

  #41683·copybara-service[bot] opened Oct 16, 2026Bot·

  ·
  

• Support non-sum UnreducedKind across JAX lax primitives.

  #41680·copybara-service[bot] opened Oct 16, 2026Bot·

  ·

• Compute log_ndtr(x) for x > 0 as log1p(-ndtr(-x)) to keep relative accuracy

  #41679·bahumanyarg11 opened Oct 16, 2026·

  ·
  1

  2

• Fixes https://github.com/jax-ml/jax/issues/41256

  #41675·copybara-service[bot] opened Oct 16, 2026Bot·

  ·
  1

• [numerics] Allow 1.5 ULP error for float32 rsqrt on non-AVX CPUs (ARM)

  #41674·copybara-service[bot] opened Oct 16, 2026Bot·

  ·
  

• [mgpu] Lower two-branch conditionals to scf.if.

  #41670·copybara-service[bot] opened Oct 16, 2026Bot·

  ·
  

• Make explicit_axes/auto_axes work with HiTypes.

  #41669·copybara-service[bot] opened Oct 16, 2026Bot·

  ·
  

• [Pallas:TPU] Add _unsafe_no_write_effect to async_copy.

  #41667·copybara-service[bot] opened Oct 16, 2026Bot·

  ·
  

• Set unique collective op_id on CPU and compute per-call Gloo tags.

  #41664·copybara-service[bot] opened Oct 16, 2026Bot·

  ·
  

• Stage core.Literal values in hijax._custom_vjp_fwd via core.stage_p.bind

  #41663·copybara-service[bot] opened Oct 15, 2026Bot·

  ·
  

• Use execution edges for collective communicator resources.

  #41661·copybara-service[bot] opened Oct 15, 2026Bot·

  ·
  1

  

• Add missing device argument to logspace and geomspace.

  pull ready
  #41659·carlosgmartin opened Oct 15, 2026Contributor·

  ·
  3

  

• [jax] Catch InternalFloatingPointError in _pjit_call_impl for debug_nans.

  #41655·copybara-service[bot] opened Oct 15, 2026Bot·

  ·
  

• [jax] Support symbolic_zeros=True in linear_transpose and use it in hijax transpose rules.

  #41654·copybara-service[bot] opened Oct 15, 2026Bot·

  ·
  

• Preserve explicit and unreduced sharding in lax.squeeze, lax.stack, and lax.unstack.

  #41653·copybara-service[bot] opened Oct 15, 2026Bot·

  ·

• Add regression test for chlo.zeta bug fix

  #41652·copybara-service[bot] opened Oct 15, 2026Bot·

  ·
  

• [mosaic_gpu] Added a dialect lowering rule for AsyncLoadIm2colOp

  #41650·copybara-service[bot] opened Oct 15, 2026Bot·

  ·
  

• nn: stabilize squareplus values and derivatives

  #41649·Robertboy18 opened Oct 15, 2026·

  ·
  1

• Default lax.empty to current mesh sharding and register zero JVP rules for empty primitives.

  #41648·copybara-service[bot] opened Oct 15, 2026Bot·

  ·

• Patch XLA with https://github.com/openxla/xla/pull/50820 (Update XPU Triton hook).

  #41643·copybara-service[bot] opened Oct 15, 2026Bot·

  ·