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:
14,301 Diff Delta
Contributors:
27 total committers
Data confidence:
Commits:
188 impacting commits

188 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
[pyrefly] add type checking for jax/experimental/sparse
Authored March 10, 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
[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
[pyrefly] type-check jax.experimental/array_serialization & colocated_python
Authored March 10, 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
Make main process_*() arguments position-only
Authored March 2, 2026
[Mosaic] NFC: Refactor tpu::PackMaskOp to take > 2 sources
Authored March 5, 2026
[pyrefly] fix errors in jax/interpreters/ad
Authored February 27, 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
[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
Bugfix for colocated python with `use_weakrefs_at_backend=True`
Authored March 12, 2026
[pyrefly] fix errors in jax/_src/shard_map.py
Authored March 2, 2026
Tracer is now generic parameterized by Trace
Authored March 13, 2026
Tie LinearizeTrace.tag to tangent_trace.tag
Authored February 26, 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
11,083 Ξ”
jax-v0.9.2 Released March 18, 2026
14,301 Ξ”
jax-v0.9.1 Released February 28, 2026
17,997 Ξ”
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

We haven't yet finished calculating and confirming the files and directories changed in this release. Please check back soon.

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(...)).
GitClear uses cookies to ensure you get the best experience on our website. Learn more about our cookie policy