Search results

<code>jax.nn.standardize</code> returns <code>nan</code> when variance is lower than <code>-epsilon</code>

bugSomething isn't workingSomething isn't working

Status: Open.

#30426 In jax-ml/jax;

· bondquant opened on Jul 23, 2025

Not Implemented functions in Pallas

bugSomething isn't workingSomething isn't working

Status: Open.

#30423 In jax-ml/jax;

· cdelv opened on Jul 23, 2025

Revisiting <code>custom_jvp</code> with <code>pure_callback</code>

bugSomething isn't workingSomething isn't working

Status: Open.

#30401 In jax-ml/jax;

· bwohlberg opened on Jul 22, 2025

Pallas fails to write to output ref when using bfloat16 on TPU v3-8

bugSomething isn't workingSomething isn't working

Status: Open.

#30391 In jax-ml/jax;

· sr5434 opened on Jul 22, 2025

Profiling tool shows no GPU processes despite GPU Util being as expected.

bugSomething isn't workingSomething isn't working

Status: Open.

#30378 In jax-ml/jax;

· saschafrey opened on Jul 22, 2025

jet and equinox.nn.MLP

bugSomething isn't workingSomething isn't working

Status: Open.

#30352 In jax-ml/jax;

· A-Alaa opened on Jul 21, 2025

Fusing the optimizer step into the backward pass

enhancementNew feature or requestNew feature or request

Status: Open.

#30338 In jax-ml/jax;

· martin-marek opened on Jul 19, 2025

jax.grad precision: float32 gradients of bfloat16 weights

enhancementNew feature or requestNew feature or request

Status: Open.

#30337 In jax-ml/jax;

· martin-marek opened on Jul 19, 2025

Segfault while Performing All-to-All Collective Operation on 8xH100 SXM5 (Shard + Swapaxes + Shard)

bugSomething isn't workingSomething isn't working

Status: Open.

#30335 In jax-ml/jax;

· novastar53 opened on Jul 19, 2025

Inconsistent Reduction precision in backwards computation

bugSomething isn't workingSomething isn't working

Status: Open.

#30310 In jax-ml/jax;

· sbodenstein opened on Jul 18, 2025

XlaRuntimeError: INTERNAL: Failed to capture gpu graph: an unsupported value or parameter was passed to the function

bugSomething isn't workingSomething isn't working

Status: Open.

#30297 In jax-ml/jax;

· hwr9912 opened on Jul 18, 2025

jax.numpy.sort performance regression

bugSomething isn't workingSomething isn't working

Status: Open.

#30296 In jax-ml/jax;

· drewjenks01 opened on Jul 18, 2025