Search results

[sharding-in-types] <code>jnp.reshape</code> not figuring out a reshape output sharding when it should

bugSomething isn't workingSomething isn't working

Status: Open.

#29235 In jax-ml/jax;

· PhilipVinc opened on Jun 4, 2025

custom_partitioning compatibility with lax.composite

enhancementNew feature or requestNew feature or request

Status: Open.

#29223 In jax-ml/jax;

· wenscarl opened on Jun 3, 2025

Expose ormqr, allowing more efficient linear least squares solves using QR factorization

enhancementNew feature or requestNew feature or request

Status: Open.

#29173 In jax-ml/jax;

· adconner opened on Jun 2, 2025

[sharding-in-types+shard_map] Support closing over explicit meshes when using shard map

bugSomething isn't workingSomething isn't working

Status: Open.

#29162 In jax-ml/jax;

· PhilipVinc opened on Jun 2, 2025

Fix broken navigation breadcrumbs on website

bugSomething isn't workingSomething isn't working

Status: Open.

#29155 In jax-ml/jax;

· carlosgmartin opened on Jun 1, 2025

Jax jacobian returns unexpected nan

questionQuestions for the JAX teamQuestions for the JAX team

Status: Open.

#29152 In jax-ml/jax;

· PcdPanda opened on Jun 1, 2025

Unimplemented primitive in Pallas GPU lowering: gather

bugSomething isn't workingSomething isn't working

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

Status: Open.

#29143 In jax-ml/jax;

· kylestach opened on May 30, 2025

jax reported a 'Segmentation fault' error for backend_compile in compiler.py

bugSomething isn't workingSomething isn't working

Status: Open.

#29139 In jax-ml/jax;

· Yunkai-Chen opened on May 30, 2025

It works fine on a single GPU, but does not function properly in a multi-GPU (parallel GPU) setup.

bugSomething isn't workingSomething isn't working

Status: Open.

#29088 In jax-ml/jax;

· juju0111 opened on May 29, 2025

Runtime failure/regression while running inference on JAX 0.6.1

bugSomething isn't workingSomething isn't working

Status: Open.

#29070 In jax-ml/jax;

· Rohanjames1997 opened on May 28, 2025

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

bugSomething isn't workingSomething isn't working

Status: Open.

#29068 In jax-ml/jax;

· galah92 opened on May 28, 2025

lax.cond crashes on Windows

bugSomething isn't workingSomething isn't working

Status: Open.

#29049 In jax-ml/jax;

· lnabergall opened on May 28, 2025