Search results

Sparse cotangent decorator

enhancementNew feature or requestNew feature or request

Status: Open.

#36862 In jax-ml/jax;

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

lax.scan with float64 3×3 matrix carry is ~10× slower on CPU in JAX 0.4.33 vs 0.4.31 (float32 unaffected)

bugSomething isn't workingSomething isn't working

Status: Open.

#36799 In jax-ml/jax;

· olivialynn opened on Apr 14, 2026

L-BFGS: improve robustness with Moré-Thuente line search and numerical safeguards

Status: Open.

#36793 In jax-ml/jax;

· jjmorenog620 opened on Apr 14, 2026

Error when the same jax.vjp pullback function is returned from the true and false functions of jax.lax.cond

bugSomething isn't workingSomething isn't working

Status: Open.

#36776 In jax-ml/jax;

· kevinwang27 opened on Apr 14, 2026

jnp.searchsorted fails with symbolic shapes on the sorted dimension during jax.export

bugSomething isn't workingSomething isn't working

Status: Open.

#36766 In jax-ml/jax;

· n-gao opened on Apr 14, 2026

Incorrect output of jax.scipy.expon.ppf for scale != 1

bugSomething isn't workingSomething isn't working

Status: Open.

#36757 In jax-ml/jax;

· solano opened on Apr 14, 2026

TPU MLIR lowering fails with illegal tpu.reshape on TPU7x

Status: Open.

#36750 In jax-ml/jax;

· shungcp opened on Apr 14, 2026

XLA topk-decomposer crash in pmap(vmap(fn)) with top_k[0] in jnp.where + separate top_k[1]

type:Bug

Status: Open.

#36703 In jax-ml/jax;

· jz-sa opened on Apr 13, 2026

jax.export.deserialize().call() fails with internal XLA error for shard_map across 2x2 mesh

bugSomething isn't workingSomething isn't working

Status: Open.

#36691 In jax-ml/jax;

· nickmcgreivy opened on Apr 11, 2026

D2H memcpy is followed by a Waiting for Host Callback Execution stall of 16 ms

bugSomething isn't workingSomething isn't working

Status: Open.

#36651 In jax-ml/jax;

· JulesGuillou opened on Apr 9, 2026