jax · Issues· 2577 open

[Pallas TPU] Support SparseCore-to-TensorCore Remote DMA on Cloud TPU

Failed build: CI - with Numpy/Scipy nightly wheels (nightly)

`jnp.concatenate` on minor axes significantly slower in newer versions

pure_callback: cross-platform sharding loses its own error message (except IndexError should be ValueError)

Support profiling without source code changes

Per-operator (non-jitted) kernels return wrong values / NaNs on RTX A6000 (CUDA 13), jitted path and CPU are correct; breaks a jitted training loop

XLA change significantly increase the computational time

[Pallas TPU] `jax.named_scope` splits otherwise co-issued instructions into separate bundles

`jnp.arcsinh` / `jnp.arccosh` JVP is exactly 0 in float32 for $|x| > ~1.84e19$ because of squared x

Improve precision of 32-bit `gammaln`?

sph_harm_y returns incorrect values for batched angles (same root cause as #20769)

GPU detected successfully but first CUDA allocation fails with CUDA_ERROR_UNKNOWN on WSL2 (RTX A4000)

betaln(a, b) is too large by exactly min(a, b) when min(a, b)/max(a, b) is subnormal

jax.scipy.special.spence returns nan over the top two binades of float32 and float64

[jit] jax_high_dynamic_range_gumbel is missing from the JIT cache key