Search results

[lax.map] Map over multiple arguments

enhancementNew feature or requestNew feature or request

Status: Open.

#29538 In jax-ml/jax;

· PhilipVinc opened on Jun 18, 2025

[regression] JAX 0.6.2 breaks <code>xp.exp(x * xp.nan)</code> inside jax.jit

bugSomething isn't workingSomething isn't working

Status: Open.

#29537 In jax-ml/jax;

· crusaderky opened on Jun 18, 2025

ValueError: Mosaic failed to compile TPU kernel: Not Implemented: The last dim size is not 128 in original base memref

bugSomething isn't workingSomething isn't working

Status: Open.

#29520 In jax-ml/jax;

· karajendran opened on Jun 17, 2025

//tests/mosaic:gpu_test_distributed_gpu and //tests/pallas:gpu_pallas_distributed_test_gpu fail on single-GPU machines

bugSomething isn't workingSomething isn't working

Status: Open.

#29513 In jax-ml/jax;

· olupton opened on Jun 17, 2025

CPU Over-utilization and taskset

bugSomething isn't workingSomething isn't working

Status: Open.

#29499 In jax-ml/jax;

· markus7800 opened on Jun 16, 2025

Unimplemented primitive in Pallas TPU lowering for KernelType.TC: dynamic_slice.

enhancementNew feature or requestNew feature or request

Status: Open.

#29481 In jax-ml/jax;

· KissShane opened on Jun 16, 2025

cuSolver internal error in <code>jnp.linalg.eigvalsh(jnp.eye(2))</code>

bugSomething isn't workingSomething isn't working

Status: Open.

#29475 In jax-ml/jax;

· Vilin97 opened on Jun 14, 2025

Missing <code>cost_estimate</code> in MHA splash attention kernel

pallasIssues pertaining to Pallas (GPU or TPU)Issues pertaining to Pallas (GPU or TPU)

Status: Open.

Bug

#29463 In jax-ml/jax;

· ZacCranko opened on Jun 13, 2025

Show uv installation instructions in the docs

enhancementNew feature or requestNew feature or request

Status: Open.

#29451 In jax-ml/jax;

· 3f6a opened on Jun 13, 2025

Transposed preconditioned GMRES

bugSomething isn't workingSomething isn't working

Status: Open.

#29449 In jax-ml/jax;

· f0uriest opened on Jun 13, 2025

<code>jax.scipy.fft.idctn</code> brings different results with <code>scipy.fft.idctn</code>

bugSomething isn't workingSomething isn't working

Status: Open.

#29426 In jax-ml/jax;

· apiqwe opened on Jun 12, 2025

Experimental support for AMD GPUs on WSL2

enhancementNew feature or requestNew feature or request

Status: Open.

#29400 In jax-ml/jax;

· Artoriuz opened on Jun 11, 2025