Search results

Add clarification to softmax and log_softmax regarding masked-out elements.

kokoro:force-run

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

Status: Open (in progress).

jax-ml/jaxnumber 30167

#30167 In jax-ml/jax;

· carlosgmartin opened on Jul 11, 2025

[direct-linearize] fix hijax tests with JAX_USE_DIRECT_LINEARIZE=1

kokoro:force-run

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

Status: Open (in progress).

jax-ml/jaxnumber 30161

#30161 In jax-ml/jax;

· mattjj opened on Jul 11, 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

[mutable-arrays] remat discharge rule

kokoro:force-run

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

Status: Open (in progress).

jax-ml/jaxnumber 29370

#29370 In jax-ml/jax;

· mattjj opened on Jun 10, 2025

[ROCm] ROCm7 Plugin Updates

kokoro:force-run

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

Status: Open (in progress).

jax-ml/jaxnumber 29281

#29281 In jax-ml/jax;

· mrodden opened on Jun 5, 2025