Search results

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

TPU Invalid relayout: Non-singleton logical dimension is replicated in destination but not in source

bugSomething isn't workingSomething isn't working

Status: Open.

#36970 In jax-ml/jax;

· AmesingFlank opened on Apr 19, 2026

Unexpected result involving scan and vmap

bugSomething isn't workingSomething isn't working

Status: Open.

#36958 In jax-ml/jax;

· puct9 opened on Apr 18, 2026

Performance regression in batched jnp.linalg.solve on CPU from jax/jaxlib 0.9.2 to 0.10.0

bugSomething isn't workingSomething isn't working

Status: Open.

#36927 In jax-ml/jax;

· clemisch opened on Apr 17, 2026

jax._src.prng.PRNGKeyArray should be consistent as jax array in truthiness evaluation

bugSomething isn't workingSomething isn't working

Status: Open.

#36907 In jax-ml/jax;

· ChromeHearts opened on Apr 16, 2026

System freeze on PREEMPT_RT during simple addition loop

bugSomething isn't workingSomething isn't working

Status: Open.

#36906 In jax-ml/jax;

· keunjun92 opened on Apr 16, 2026

shard_map + grad produces unexpected results on multi-device

bugSomething isn't workingSomething isn't working

Status: Open.

#36886 In jax-ml/jax;

· msolinas28 opened on Apr 16, 2026

Feature request: exploit symmetric structure in matrix-vector/matrix-matrix products (optimal bilinear complexity)

enhancementNew feature or requestNew feature or request

Status: Open.

#36859 In jax-ml/jax;

· MahmoodKhalil57 opened on Apr 16, 2026

Unexpected Numerical Error Trend in jax.vmap with jax.lax.conv_general_dilated on TPU

bugSomething isn't workingSomething isn't working

Status: Open.

#36813 In jax-ml/jax;

· BensonHuangTW opened on Apr 15, 2026