Search results

[Mosaic GPU] Fix some typos in docs

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 30447

#30447 In jax-ml/jax;

· justinjfu opened on Jul 23, 2025

Skip Pallas and Mosaic GPU tests that don't fit on RTX 6000 PRO

CI Optional GPU PresubmitLabel to flag PR to run additional GPU testing not in standard presubmitsLabel to flag PR to run additional GPU testing not in standard presubmits

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 30258

#30258 In jax-ml/jax;

· andportnoy opened on Jul 16, 2025

[jaxprs] Hoist large constants as arguments during lowering

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 30180

#30180 In jax-ml/jax;

· gnecula opened on Jul 14, 2025

Document all checkpoint policies in one place, on the JAX public API page.

documentation

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 30177

#30177 In jax-ml/jax;

· carlosgmartin opened on Jul 12, 2025

Fix cache leaks for pe._cached_abstract_eval; add util.multi_weakref_lru_cache

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 30009

#30009 In jax-ml/jax;

· gnecula opened on Jul 7, 2025

Add array-api copy semantics to DLPackManagedTensorToBuffer

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 29963

#29963 In jax-ml/jax;

· pearu opened on Jul 3, 2025

Add nd_loop and Enable block_n tiling for all_gather_lhs_matmul

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 29822

#29822 In jax-ml/jax;

· hanzlfs opened on Jun 27, 2025

[CI] Add bazel TPU presubmit testing

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 29660

#29660 In jax-ml/jax;

· MichaelHudgins opened on Jun 23, 2025

Add support for unroll=0 to lax.scan.

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 29630

#29630 In jax-ml/jax;

· carlosgmartin opened on Jun 21, 2025

[mutable-arrays] enable checks by default

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 29620

#29620 In jax-ml/jax;

· mattjj opened on Jun 20, 2025

Add <code>split_p</code> to <code>jax.extend.core.primitives</code>

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 29452

#29452 In jax-ml/jax;

· j-towns opened on Jun 13, 2025

Update missed ml_dtypes version check to 0.5.0

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 29429

#29429 In jax-ml/jax;

· DanisNone opened on Jun 12, 2025