Search results

jax.nn.standardize produces different outputs in eager vs. outer-JIT context due to numerically unstable variance formula

bugSomething isn't workingSomething isn't working

Status: Open.

#37281 In jax-ml/jax;

· jasminetrail opened on Apr 29, 2026

scipy.fft dst and idst type 2

enhancementNew feature or requestNew feature or request

Status: Open.

#37268 In jax-ml/jax;

· Nin17 opened on Apr 29, 2026

pl.kernel returns mutable Ref out of JIT under jax.lax.scan + jax.grad

Status: Open.

#37251 In jax-ml/jax;

· ultrons opened on Apr 28, 2026

jax.jit produces NaN/Inf from softmax when pre-activation values are very large; eager returns finite output

bugSomething isn't workingSomething isn't working

Status: Open.

#37234 In jax-ml/jax;

· jasminetrail opened on Apr 28, 2026

TPU v6e init fails: TPU_RET_CHECK GetChip(0)->location().index_on_host() == i (0 vs. 1) (jax 0.9.2)

Status: Open.

#37230 In jax-ml/jax;

· stharrold opened on Apr 28, 2026

Input variables and their types should be carried through JitWrapped functions

enhancementNew feature or requestNew feature or request

Status: Open.

#37188 In jax-ml/jax;

· brugalada opened on Apr 27, 2026

jax.device_put silently returns zeros when multiple slurm jobs are running on the same node

bugSomething isn't workingSomething isn't working

Status: Open.

#37150 In jax-ml/jax;

· xqding opened on Apr 24, 2026

Custom lag arguments for np.convolve and np.correlate

enhancementNew feature or requestNew feature or request

Status: Open.

#37112 In jax-ml/jax;

· HoniSanders opened on Apr 23, 2026

CPU PJRT C API plugin

enhancementNew feature or requestNew feature or request

Status: Open.

#37033 In jax-ml/jax;

· markhalka opened on Apr 21, 2026

TPU IsFusibleUnalignedDUS crash on LLaMA rotate-half RoPE at head_dim=128

bugSomething isn't workingSomething isn't working

Status: Open.

#37029 In jax-ml/jax;

· pedronahum opened on Apr 21, 2026

jnp.take_along_axis fails to lower in Mosaic TPU

bugSomething isn't workingSomething isn't working

Status: Open.

#37006 In jax-ml/jax;

· AmesingFlank opened on Apr 20, 2026

Auto sharding constraint fails on multi-process jobs

bugSomething isn't workingSomething isn't working

Status: Open.

#37001 In jax-ml/jax;

· ChenAo-Phys opened on Apr 20, 2026