Home » Oreps » Jax » All Releases » jax-v0.10.1 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.10.1
Basic Details Issues Addressed Top Contributors Directory Browser Release Notes
Jax jax-v0.10.1

Jax: jax-v0.10.1 Release

Release date:
May 18, 2026
Previous version:
jax-v0.10.0 (released April 16, 2026)
Magnitude:
31,116 Diff Delta
Contributors:
53 total committers
Data confidence:
Commits:
489 impacting commits

489 Commits in this Release

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

Make stack and unstack jax primitives
Authored May 12, 2026
Add a jaxpr_html IR dump mode that dumps a jaxpr to an HTML document
Authored May 4, 2026
[docs] use jit decorator factory throughout
Authored May 7, 2026
Delete pjit.AUTO support from JAX
Authored April 23, 2026
Move PartitionSpec back to python and intern it
Authored April 28, 2026
Create `pl.select_ref` to dynamically choose from multiple refs for DMAs
Authored May 5, 2026
Define Tracer.dimension_as_value and Tracer.__jax_array__ as None
Authored April 16, 2026
Simplify the flattening of TransformedRefs in DMAs
Authored April 20, 2026
Move PartitionSpec back to python and intern it
Authored April 28, 2026
[Pallas][TPU kernel interpreter] Add an interpreter option that enables padding allocations
Authored April 27, 2026
Fix pallas fuser BlockIndexTransform comparison to consider symbolic equivalence of block indices transforms rather than just object identity
Authored May 7, 2026
replace fusible_p with Fusible(VJPHiPrimitive)
Authored April 22, 2026
[gh] Post-process workflow test results and load into BigQuery
Authored May 6, 2026
[JAX] Cache the aval_to_ir_type mapping during MLIR lowering
Authored April 30, 2026
[Pallas][GPU kernel interpreter] Clarify type annotations, esp
Authored May 8, 2026
jax: replace trace_to_jaxpr_dynamic with trace_to_jaxpr in state and interpreters
Authored May 15, 2026
[Pallas] Add core_map tests to tpu_fusible_matmul_test
Authored April 20, 2026
[XLA:MGPU] Delete C++ implementation of core primitives
Authored April 27, 2026
Roll back https://github.com/jax-ml/jax/pull/37117
Authored May 4, 2026
[Mosaic GPU] Reject invalid assignments when reducing constraint system
Authored April 21, 2026
[Mosaic GPU] Use `fa.plan_tiled_transfer` in order to validate layouts in `IsTransferableSmemRegisters` constraints
Authored April 16, 2026
Use Pyrefly suppressions instead of `pytype: disable=`
Authored April 15, 2026
PR #37117: jax/mosaic: enable CUPTI V2 multi-subscriber coexistence on CUDA 13.2+
Authored May 4, 2026
Don't use nanobind to construct Doc objects
Authored May 1, 2026
[Pallas][GPU kernel interpreter] Support simulating multiple GPUs with `shard_map`
Authored May 8, 2026
Removed the now-obsolete workarounds for pytype
Authored April 16, 2026
jax/lax: replace trace_to_jaxpr_dynamic with trace_to_jaxpr
Authored May 14, 2026
Replace compute_on tests with compute_on2
Authored May 14, 2026
[Mosaic GPU] Generate multiple pairwise PTX conversions when possible
Authored May 8, 2026
Remove sharding_impls.AxisEnv
Authored April 29, 2026
[Mosaic GPU] Allow collapsing unit dimensions when issuing TMA loads
Authored April 28, 2026
[Mosaic:GPU] Re-land use NCCL device API based barrier to support cross-host barrier on a single cluster
Authored April 30, 2026
[Mosaic GPU] Support multidimensional transposes in layout inference
Authored May 4, 2026
Add HTML color support to the pretty printer
Authored May 7, 2026
Reverts 6fe3e4626169a4bb30cb87526599f5d06b7c467a
Authored May 12, 2026
[Pallas/Mosaic GPU] Implement logic that allows commuting `ReshapeTransform`s with `UntilingTransform`s
Authored April 28, 2026
[Jax] Refactor `jax.debug.callback` and `jax.debug.print` to support curried options
Authored April 23, 2026
Make SdyArray immutable and cache SdyArray to sharding attribute conversion
Authored April 21, 2026
[pallas:sc] Added `plsc.Indices` which allow specifying the ignored offset value for gathers/scatters
Authored April 28, 2026
[Pallas:MGPU] Add Warpgroup lowering for async_store_smem
Authored May 1, 2026

Browse Other Releases

Latest Pending Unreleased 😎
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 May 15, 2026
429,407 Ξ”
jax-v0.10.1 Released May 18, 2026
31,116 Ξ”
jax-v0.10.0 Released April 16, 2026
62,264 Ξ”
jax-v0.9.2 Released March 18, 2026
14,078 Ξ”
jax-v0.9.1 Released February 28, 2026
17,693 Ξ”
jax-v0.9.0.1 Released February 3, 2026
0 Ξ”
jax-v0.9.0 Released January 20, 2026
0 Ξ”

Top Contributors in jax-v0.10.1

hawkinsp
a-googler
yashk2810
superbobry
bchetioui
apaszke
golechwierowicz
allanrenucci
jakevdp
IvyZX

Directory Browser for jax-v0.10.1

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

Loading File Browser...

Release Notes Published

  • New features

    • Added ResizeMethod.AREA to jax.image.resize, which matches TensorFlow's AREA resizing (#20098).
    • Added jax.scipy.linalg.hadamard for constructing Hadamard matrices (#10144).
    • Added jax.scipy.linalg.circulant for constructing circulant matrices (#10144).
    • Added jax.scipy.linalg.dft for constructing discrete Fourier transform matrices (#10144).
    • Added jax.scipy.linalg.leslie for constructing Leslie matrices (#10144).
    • Added jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients (#10144).
    • Added jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices (#10144).
    • Added jax.scipy.linalg.helmert for constructing Helmert matrices (#10144).
    • Moved RNG APIs from "implementations" to dtypes (#27854):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

  • Breaking changes

    • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

  • Deprecations

    • Passing the copy, order, and ndmin arguments to jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
    • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future version of JAX, this will become an error, if you depend on using them as leaves, pass is_leaf to jax.tree.* methods.