Home » Oreps » Jax » All Releases » jax-v0.9.2 Release
google/jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
Watch
1
Commit Activity
Directory Browser
Velocity History
Releasesjax-v0.9.2
Basic Details Issues Addressed Top Contributors Directory Browser Release Notes
Jax jax-v0.9.2

Jax: jax-v0.9.2 Release

Release date:
March 18, 2026
Previous version:
jax-v0.9.1 (released February 28, 2026)
Magnitude:
15,566 Diff Delta
Contributors:
32 total committers
Data confidence:
Commits:
195 impacting commits

195 Commits in this Release

Ordered by the degree to which they evolved the repo in this version.

internal change
Authored March 12, 2026
[Take 2] shard map without linear util
Authored March 13, 2026
[0.9.1] Update jaxlib version guards post-release
Authored March 4, 2026
[Pallas/TPU] Add tpu_info.get_tpu_info_for_chip
Authored February 26, 2026
[pyrefly] type check jax/experimental/pallas
Authored March 6, 2026
Add workflow for build wheels and running bazel tests on ROCm
Authored January 21, 2026
[Pallas][TPU kernel interpreter] Add optional logging to memory operations
Authored March 4, 2026
[Mosaic GPU] Add support for SMEM aliasing using WG semantics in the Pallas lowering
Authored March 5, 2026
[pyrefly] add type checking for jax/experimental/sparse
Authored March 10, 2026
[Pallas][TPU kernel interpreter] Consistently use `pltu.HBM` to refer to HBM in the interpreter
Authored February 27, 2026
[NFC] Remove and modify some out-of-date test skips
Authored March 2, 2026
Mosaic GPU: pre-launch check for multicast support
Authored February 10, 2026
Improve safe_zip error messages to show argument lengths
Authored March 5, 2026
[pyrefly] fix pyrefly errors in jax._src.core
Authored March 2, 2026
[Mosaic:GPU] Re-use the barrier buffers between executions to prevent situation when the buffer were released before the end of cross-device barrier
Authored February 27, 2026
[Pallas][Mosaic TPU] More complete 1D block support in Pallas/Mosaic
Authored February 28, 2026
[export] Remove code to serialize HloSharding
Authored March 3, 2026
[pyrefly] type-check jax.experimental/array_serialization & colocated_python
Authored March 10, 2026
[hijax] cleaner scan inc_rank/dec_rank api
Authored February 26, 2026
Add `HloShardingV3 -> JAX NamedSharding` conversion, used for shardings returned via compiled XLA executable back to JAX
Authored March 9, 2026
Show detailed mismatch key paths in `flatten_axes()` tree prefix errors
Authored March 7, 2026
[typing] improve annotations for jax/_src/test_util
Authored March 5, 2026
[refactor] directly define array methods
Authored March 4, 2026
[Pallas/Mosaic GPU] Turn on `gpu_pallas_distributed_test`s for warpgroup semantics
Authored March 13, 2026
Simplify input validation for jax.Array operator overloads
Authored March 5, 2026
scipy.linalg: add hankel special matrix
Authored February 19, 2026
primitives: update output dtype handling for dot function. 0
Authored March 5, 2026
Make main process_*() arguments position-only
Authored March 2, 2026
[Mosaic] NFC: Refactor tpu::PackMaskOp to take > 2 sources
Authored March 5, 2026
[maint] clean up array class overrides
Authored March 12, 2026
[Mosaic GPU] Refactor `FragmentedArray.broadcast` to use `broadcast_in_dim`
Authored March 4, 2026
[typing] miscellaneous pyrefly typing fixes
Authored February 26, 2026
[Mosaic GPU][NFC] Remove dead `constraints.BroadcastInDim`
Authored March 2, 2026
Add C addition to hopper matmul kernel
Authored March 2, 2026
[pyrefly] fix errors in jax/interpreters/ad
Authored February 27, 2026
[doc] update AI code policy
Authored March 9, 2026
[Mosaic GPU]Β Add `a_sparse_metadata` argument to `tcgen05_mma` dialect definition
Authored March 6, 2026
Reverts ca6fda28429a4f4a17caeb95e4e3f45a14669b01
Authored March 5, 2026
[Mosaic] Simplify GetCoreTypeOfParentOp
Authored February 28, 2026
[Pallas/Mosaic GPU] Enable lowering `semaphore_{read,signal,wait}_p` using warpgroup semantics
Authored March 10, 2026

Browse Other Releases

Latest Pending Unreleased 😎
jaxlib-v0.1.39 Released February 14, 2020
0 Ξ”
jaxlib-v0.1.38 Released January 29, 2020
0 Ξ”
jaxlib-v0.1.37 Released December 6, 2019
0 Ξ”
jaxlib-v0.1.36 Released November 22, 2019
0 Ξ”
jaxlib-v0.1.33 Released January 21, 2026
12,194 Ξ”
jax-v0.9.2 Released March 18, 2026
15,566 Ξ”
jax-v0.9.1 Released February 28, 2026
19,549 Ξ”
jax-v0.9.0.1 Released February 3, 2026
0 Ξ”
jax-v0.9.0 Released January 20, 2026
0 Ξ”
jax-v0.8.3 Released January 28, 2026
28 Ξ”

Top Contributors in jax-v0.9.2

a-googler
jakevdp
danielsuo
dougalm
yueshengys
bchetioui
allanrenucci
WindQAQ
dimitar-asenov
charleshofer

Directory Browser for jax-v0.9.2

All files are compared to previous version, jax-v0.9.1. Click here to browse diffs between other versions.

Loading File Browser...

Release Notes Published

JAX 0.9.2 (March 2, 2026)

  • Changes:
    • The semi-private type jax._src.literals.TypedNdArray is now a subclass of np.ndarray, rather than a duck type of it.
    • jax.numpy.arange with step specified no longer generates the array on host. The benefit is more efficient code, though this can lead to less precise outputs for narrow-width floats (e.g. bfloat16). To recover the previous behavior in this case, use jnp.array(np.arange(...)).