Search results

jit C++ fastpath attaches the logical instead of the physical sharding to PRNG key outputs, causing a spurious cache miss when keys are fed back in

bugSomething isn't workingSomething isn't working

Status: Open.

#39089 In jax-ml/jax;

· Gattocrucco opened on Jul 11, 2026

CPU FFT becomes much slower when combined with adjacent operations in one jax.jit

bugSomething isn't workingSomething isn't working

Status: Open.

#39081 In jax-ml/jax;

· gRox167 opened on Jul 11, 2026

Full explicit-attributes C++ pjit cache (8192 live entries) permanently stops fastpath retention for all later jits, plain jit included

Status: Open.

#39078 In jax-ml/jax;

· subho406 opened on Jul 10, 2026

Pallas TPU: document unpack_elementwise uint32-source + block-tiling requirements

Status: Open.

#39058 In jax-ml/jax;

· lokic233 opened on Jul 10, 2026

Silent wrong gradient: eager (non-jit) VJP of x[:, idx, :] gather zeroes batch rows on TPU (regression in libtpu 0.0.42)

bugSomething isn't workingSomething isn't working

TPUIssues related to compilation for or performance on Google TPUsIssues related to compilation for or performance on Google TPUs

XLA

Status: Open.

#39025 In jax-ml/jax;

· m-braganca opened on Jul 9, 2026

Support more callable functions for modifying hlo

enhancementNew feature or requestNew feature or request

Status: Open.

#39023 In jax-ml/jax;

· Iamleos opened on Jul 9, 2026

jnp.bfloat16's smallest_subnormal truncated to zero when converted to jnp.float64 and then back to jnp.bfloat16.

bugSomething isn't workingSomething isn't working

Status: Open.

#38995 In jax-ml/jax;

· gonnet opened on Jul 8, 2026

Pallas-TPU: document/support vmem_limit_bytes > 64MB for fused dequant kernels (tpu7x has more)

Status: Open.

#38988 In jax-ml/jax;

· lokic233 opened on Jul 8, 2026

Pallas-TPU: vector.extract_strided_slice stride verification error (blocks pre-split weight unpack)

Status: Open.

#38987 In jax-ml/jax;

· lokic233 opened on Jul 8, 2026

Pallas-TPU: arith.shrui on int8 vectors fails Mosaic legalization (blocks in-kernel nibble unpack)

Status: Open.

#38986 In jax-ml/jax;

· lokic233 opened on Jul 8, 2026

_SubParsersAction.__init__() got an unexpected keyword argument 'required'

bugSomething isn't workingSomething isn't working

Status: Open.

#38969 In jax-ml/jax;

· kiefernforst opened on Jul 7, 2026

Document io_callback with sharding

enhancementNew feature or requestNew feature or request

Status: Open.

#38968 In jax-ml/jax;

· Findus23 opened on Jul 7, 2026