Search results

Incorrect gradient of lax.digamma for negative half-integer inputs

bugSomething isn't workingSomething isn't working

Status: Open.

#37635 In jax-ml/jax;

· ALinrunrun opened on May 13, 2026

Incorrect jacfwd result for logsumexp with zero b coefficient

bugSomething isn't workingSomething isn't working

Status: Open.

#37633 In jax-ml/jax;

· ALinrunrun opened on May 13, 2026

Potential issue with shard_map + ordered IO callbacks that throw exceptions

bugSomething isn't workingSomething isn't working

Status: Open.

#37587 In jax-ml/jax;

· jburnim opened on May 12, 2026

CPU compile breaks JAX profiler writes TPU HloProtos to /host:metadata

bugSomething isn't workingSomething isn't working

Status: Open.

#37533 In jax-ml/jax;

· jaustin-ant opened on May 11, 2026

# [Perf] XLA fails to fuse segmented reductions in softmax (segment_max → exp → segment_sum), 8-19x slower than hand-fused Triton

bugSomething isn't workingSomething isn't working

Status: Open.

#37511 In jax-ml/jax;

· Umang-projects opened on May 8, 2026

CPU performance regression

bugSomething isn't workingSomething isn't working

Status: Open.

#37465 In jax-ml/jax;

· clemisch opened on May 7, 2026

Closing over numpy array constant with JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS=True crashes while lowering

bugSomething isn't workingSomething isn't working

Status: Open.

#37445 In jax-ml/jax;

· bzh-bzh opened on May 7, 2026

jaxlib: malformed pickled PyTreeDef can crash flatten_up_to

bugSomething isn't workingSomething isn't working

Status: Open.

#37412 In jax-ml/jax;

· sjfbo opened on May 6, 2026

jaxlib: malformed PyTreeDefProto can crash process in PyTreeDef.deserialize_using_proto

bugSomething isn't workingSomething isn't working

Status: Open.

#37410 In jax-ml/jax;

· sjfbo opened on May 6, 2026

jax.config.update("jax_compiler_enable_remat_pass", False) does not respect existing XLA_FLAGS

bugSomething isn't workingSomething isn't working

Status: Open.

#37391 In jax-ml/jax;

· SobhanMP opened on May 5, 2026

jax.custom_jvp produces non-deterministically wrong results (exact +2 offset) when called via vmap after processing other custom-pytree types in the same process

bugSomething isn't workingSomething isn't working

Status: Open.

#37389 In jax-ml/jax;

· NeilGirdhar opened on May 5, 2026

[Apple GPU (Metal) plugin] Deadlock in jax-metal at block_until_ready

bugSomething isn't workingSomething isn't working

Status: Open.

#37374 In jax-ml/jax;

· wolf-openai opened on May 4, 2026