Search results

[Feature Request] MCP Server for AI-Powered Model Development

Status: Open.

#34212 In jax-ml/jax;

· PurpleSquirrelMedia opened on Jan 7, 2026

Using jax.make_jaxpr with <code>return_shape=True</code> on a functions that return a HiType fail

bugSomething isn't workingSomething isn't working

Status: Open.

#34193 In jax-ml/jax;

· armandpicard opened on Jan 7, 2026

have ROCm test parity with CUDA &amp; TPUs in upstream

enhancementNew feature or requestNew feature or request

Status: Open.

#34184 In jax-ml/jax;

· functionstackx opened on Jan 7, 2026

Expert parallelism and "Unimplemented group_offset support" for jax.lax.ragged_dot

enhancementNew feature or requestNew feature or request

Status: Open.

#34168 In jax-ml/jax;

· pcmoritz opened on Jan 6, 2026

<code>jnp.sinc</code>'s gradient has large errors

bugSomething isn't workingSomething isn't working

Status: Open.

#34139 In jax-ml/jax;

· howsiyu opened on Jan 3, 2026

Silent failure when indexing using tuples

questionQuestions for the JAX teamQuestions for the JAX team

Status: Open.

#34127 In jax-ml/jax;

· SNMS95 opened on Dec 31, 2025

Improve __repr__ for jax.nn functions and initializers

enhancementNew feature or requestNew feature or request

Status: Open.

#34124 In jax-ml/jax;

· hamzamerzic opened on Dec 30, 2025

[TPU][Mosaic] The mosaic IR dump is using a filename that is too large

bugSomething isn't workingSomething isn't working

Status: Open.

#34117 In jax-ml/jax;

· patrick-toulme opened on Dec 29, 2025

jax-metal fails to run on recent macOS

Apple GPU (Metal) plugin

bugSomething isn't workingSomething isn't working

Status: Open.

#34109 In jax-ml/jax;

· FlorinAndrei opened on Dec 26, 2025

<code>jax.nn.dot_product_attention</code> creates internal INFs when using float64 data

bugSomething isn't workingSomething isn't working

Status: Open.

#34107 In jax-ml/jax;

· burnpanck opened on Dec 26, 2025

Regression in JAX 0.8.2: XLA <code>permutation_sort_simplifier</code> HLO verifier crashes compiling argsort-based inverse permutation (int64/x64)

bugSomething isn't workingSomething isn't working

Status: Open.

#34096 In jax-ml/jax;

· bmucsanyi opened on Dec 23, 2025

XLA GPU: “Contracting dimension is too fragmented” triggered by certain seq_len × batch combinations

bugSomething isn't workingSomething isn't working

Status: Open.

#34094 In jax-ml/jax;

· lych1233 opened on Dec 23, 2025