Search results

[Pallas TPU] emit_pipeline fails on SparseCore without shard_map

bugSomething isn't workingSomething isn't working

Status: Open.

#34640 In jax-ml/jax;

· thorsley1 opened on Jan 26, 2026

Weird tensor shape when computing the gradient of batched grouped convolutions with a custom VJP

bugSomething isn't workingSomething isn't working

Status: Open.

#34639 In jax-ml/jax;

· LambdaP opened on Jan 26, 2026

It is not possible to define the .vma of FFI outputs

bugSomething isn't workingSomething isn't working

Status: Open.

#34626 In jax-ml/jax;

· jstuecker opened on Jan 25, 2026

Convenient interface for raising exceptions via a conditional host-callback

enhancementNew feature or requestNew feature or request

Status: Open.

#34624 In jax-ml/jax;

· jstuecker opened on Jan 25, 2026

[Pallas TPU] jnp.argmin/argmax returns last index instead of first index on ties

bugSomething isn't workingSomething isn't working

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

Status: Open.

#34620 In jax-ml/jax;

· catswe opened on Jan 24, 2026

Unsupported ROCm Tests <code>tests/scaled_matmul_stablehlo_test.py</code>

AMD GPUIssues pertaining to AMD GPUs (ROCM)Issues pertaining to AMD GPUs (ROCM)

bugSomething isn't workingSomething isn't working

Status: Open.

#34519 In jax-ml/jax;

· JehandadKhan opened on Jan 21, 2026

Unsupported ROCm Tests <code>tests/pallas/triton_pallas_test.py</code>

AMD GPUIssues pertaining to AMD GPUs (ROCM)Issues pertaining to AMD GPUs (ROCM)

bugSomething isn't workingSomething isn't working

Status: Open.

#34518 In jax-ml/jax;

· JehandadKhan opened on Jan 21, 2026

Unsupported ROCm Tests <code>tests/pallas/gpu_pallas_distributed_test.py</code>

AMD GPUIssues pertaining to AMD GPUs (ROCM)Issues pertaining to AMD GPUs (ROCM)

bugSomething isn't workingSomething isn't working

Status: Open.

#34516 In jax-ml/jax;

· JehandadKhan opened on Jan 21, 2026

Unsupported ROCm Tests <code>tests/nn_test.py</code>

AMD GPUIssues pertaining to AMD GPUs (ROCM)Issues pertaining to AMD GPUs (ROCM)

bugSomething isn't workingSomething isn't working

Status: Open.

#34515 In jax-ml/jax;

· JehandadKhan opened on Jan 21, 2026

Unsupported ROCm Tests <code>tests/multiprocess_gpu_test.py</code>

AMD GPUIssues pertaining to AMD GPUs (ROCM)Issues pertaining to AMD GPUs (ROCM)

bugSomething isn't workingSomething isn't working

Status: Open.

#34514 In jax-ml/jax;

· JehandadKhan opened on Jan 21, 2026

Unsupported ROCm Tests <code>tests/lobpcg_test.py</code>

AMD GPUIssues pertaining to AMD GPUs (ROCM)Issues pertaining to AMD GPUs (ROCM)

bugSomething isn't workingSomething isn't working

Status: Open.

#34513 In jax-ml/jax;

· JehandadKhan opened on Jan 21, 2026

Unsupported ROCm Tests <code>tests/fused_attention_stablehlo_test.py</code>

AMD GPUIssues pertaining to AMD GPUs (ROCM)Issues pertaining to AMD GPUs (ROCM)

bugSomething isn't workingSomething isn't working

Status: Open.

#34512 In jax-ml/jax;

· JehandadKhan opened on Jan 21, 2026