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

Jax: jax-v0.7.1 Release

Release date:
August 19, 2025
Previous version:
jax-v0.7.0 (released July 22, 2025)
Magnitude:
19,862 Diff Delta
Contributors:
48 total committers
Data confidence:
Commits:
228 impacting commits

228 Commits in this Release

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

Refactoring HTTIJ and surrounding pages
Authored June 17, 2025
Remove jaxlib_extension_version, ifrt_version guards now that 0.7.0 release is out
Authored July 22, 2025
add training cookbook
Authored August 15, 2025
Reverts d86857a9c8cdcae8d27e6a371daceb0ac3f53e28
Authored July 22, 2025
[pallas:triton] Migrated example kernels to `plgpu.{load,store}`
Authored July 24, 2025
#sdy Delete JAX test configs enabling Shardy
Authored July 22, 2025
[Mosaic GPU]Β Use a custom pass to compile PTX to CUBIN
Authored August 18, 2025
[Mosaic GPU] Implement lowering rules for `async_tmem_load` and `async_tmem_store`
Authored July 30, 2025
Implement explicit mode sharding rule for scatter
Authored July 23, 2025
[Mosaic GPU] Add support for tcgen05 collective mma to MGPU dialect
Authored August 15, 2025
[Pallas] Update TPU docs with new pipelining features
Authored August 12, 2025
Add attention sinks to the splash attention kernel including bwd pass
Authored August 15, 2025
[Pallas][Mosaic GPU] Support pytree in/out specs in the warp specialized pipeline
Authored August 1, 2025
[Pallas][Mosaic GPU] Use separate store warpgroup for matmul kernel on Blackwell
Authored July 30, 2025
Add Hermetic C++ Toolchains for Linux x86_64 builds
Authored July 23, 2025
[Mosaic GPU] Reverse the direction in which we iterate through the pallas grid to be row-major
Authored July 25, 2025
[pallas] Forked `load` and `store` into `triton` and `tpu`
Authored July 24, 2025
[Mosaic GPU] Add initial lowering for `tcgen05_mma`
Authored August 8, 2025
Add optional dimension_numbers argument to lax.dot
Authored August 4, 2025
[Pallas TPU] Remove `use_sreg_for_state` flag
Authored July 23, 2025
Automated Code Change
Authored July 30, 2025
[mutable arrays] add API docs for jax.ref contents
Authored August 8, 2025
[Mosaic GPU][NFC] Simplify usages of `DialectBarrierRef` in MGPU dialect tests
Authored August 18, 2025
Avoid top-level imports in jax/_src/internal_test_util/*
Authored July 23, 2025
Support core axis index in the `device_id` dict for async copy and semaphore
Authored July 26, 2025
[Pallas TPU] Optimize DMA prefetch order in emit_pipeline
Authored August 6, 2025
[Mosaic GPU] Implement a rule to derive an equation system from a `vector.MultiDimReductionOp`
Authored August 15, 2025
[Mosaic GPU] Predicate `tmem_alloc` and `tmem_dealloc` at the warp level
Authored July 29, 2025
Migrate dependencies on jax experimental targets to `jax/experimental/BUILD`
Authored August 1, 2025
[Mosaic GPU] Support warpgroup semantic for automatic TMEM allocation
Authored August 5, 2025
[Pallas TPU] Expose Pallas PRNG methods to the public `pltpu` namespace and clean up tests
Authored August 14, 2025
Avoid top-level imports in jax/_src/test_util.py
Authored July 22, 2025
[pallas] Moved atomic APIs to `jax.experimental.pallas.triton`
Authored July 24, 2025
[Mosaic] Allow matrix-vector dot
Authored August 7, 2025
Extend numpy.size to accept multiple axes
Authored July 22, 2025
[Mosaic GPU] Add a `Reshape` expression constructor
Authored August 15, 2025
[Mosaic GPU] Do not call `ptxas` to query the latest supported PTX ISA
Authored August 8, 2025
[Mosaic GPU][NFC] Check that the output has the expected type in `_undo_conversion_cast`
Authored August 4, 2025
[Mosaic GPU][NFC] Move debug dumping utils into dedicated files
Authored August 15, 2025
Remove JAX_IFRT_VERSION_NUMBER guards after the release of 0.7.0
Authored July 22, 2025

Browse Other Releases

Latest Pending Unreleased 😎
jax-v0.8.2 Released December 18, 2025
22,685 Ξ”
jax-v0.8.1 Released November 18, 2025
16,240 Ξ”
jax-v0.8.0 Released October 15, 2025
18,094 Ξ”
jax-v0.7.2 Released September 15, 2025
14,732 Ξ”
jax-v0.7.1 Released August 19, 2025
19,862 Ξ”
jax-v0.7.0 Released July 22, 2025
65,989 Ξ”
jax-v0.6.2 Released June 17, 2025
39,793 Ξ”
jax-v0.6.1 Released June 18, 2025
62,649 Ξ”
jax-v0.6.0 Released April 22, 2025
67,287 Ξ”
jax-v0.5.3 Released March 31, 2025
45,270 Ξ”

Top Contributors in jax-v0.7.1

melissawm
bchetioui
allanrenucci
jakevdp
yashk2810
justinjfu
superbobry
dimitar-asenov
emilyfertig
rdyro

Directory Browser for jax-v0.7.1

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

Release Notes Published

  • New features

    • JAX now ships Python 3.14 and 3.14t wheels.
    • JAX now ships Python 3.13t and 3.14t wheels on Mac. Previously we only offered free-threading builds on Linux.

  • Changes

    • Exposed jax.set_mesh which acts as a global setter and a context manager. Removed jax.sharding.use_mesh in favor of jax.set_mesh.
    • JAX is now built using CUDA 12.9. All versions of CUDA 12.1 or newer remain supported.
    • jax.lax.dot now implements the general dot product via the optional dimension_numbers argument.

  • Deprecations:

    • jax.lax.zeros_like_array is deprecated. Please use jax.numpy.zeros_like instead.
    • Attempting to import jax.experimental.host_callback now results in a DeprecationWarning, and will result in an ImportError starting in JAX v0.8.0. Its APIs have raised NotImplementedError since JAX version 0.4.35.
    • In jax.lax.dot, passing the precision and preferred_element_type arguments by position is deprecated. Pass them by explicit keyword instead.
    • Several dozen internal APIs have been deprecated from jax.interpreters.ad, jax.interpreters.batching, and jax.interpreters.partial_eval; they are used rarely if ever outside JAX itself, and most are deprecated without any public replacement.