Search results

Try to add constants directly into core.Jaxpr

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

Status: Open (in progress).

jax-ml/jaxnumber 29582

#29582 In jax-ml/jax;

· gnecula opened on Jun 19, 2025

[debug_info] Improve debug info for loop constructs.

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

Status: Open (in progress).

jax-ml/jaxnumber 29539

#29539 In jax-ml/jax;

· gnecula opened on Jun 18, 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] re-land #29353

kokoro:force-run

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

Status: Open (in progress).

jax-ml/jaxnumber 29421

#29421 In jax-ml/jax;

· mattjj 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

Grad of unreduced

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

Status: Open (in progress).

jax-ml/jaxnumber 29219

#29219 In jax-ml/jax;

· yashk2810 opened on Jun 3, 2025

Enable more flexible handling of custom_vjp under remat by deferring partial eval

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

Status: Draft (not ready).

jax-ml/jaxnumber 29116

#29116 In jax-ml/jax;

· dfm opened on May 30, 2025

[scan] don't hoist loop-invariant computations in scan, just forward

kokoro:force-run

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

Status: Open (in progress).

jax-ml/jaxnumber 28993

#28993 In jax-ml/jax;

· mattjj opened on May 23, 2025

[JAX][DOC] Add optimizer state offloading doc

kokoro:force-run

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

Status: Open (in progress).

jax-ml/jaxnumber 28988

#28988 In jax-ml/jax;

· zhenying-liu opened on May 23, 2025

Use the builtin <code>isinstance</code> with MLIR types

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

Status: Open (in progress).

jax-ml/jaxnumber 28973

#28973 In jax-ml/jax;

· superbobry opened on May 23, 2025

Experiment: mutable RNG object

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

Status: Draft (not ready).

jax-ml/jaxnumber 28845

#28845 In jax-ml/jax;

· jakevdp opened on May 20, 2025