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

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,270 Diff Delta

Contributors:

48 total committers

Data confidence:

Commits:

228

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

Move experimental targets from `jax/BUILD` into `jax/experimental/BUILD`

Authored July 31, 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

[Mosaic GPU] Use a custom pass to compile PTX to CUBIN

Authored August 18, 2025

#sdy Delete JAX test configs enabling Shardy

Authored July 22, 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

[Mosaic GPU] Add an env variable `MOSAIC_GPU_DUMP_TO` to setup debug dumps

Authored August 6, 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][NFC] Simplify usages of `DialectBarrierRef` in MGPU dialect tests

Authored August 18, 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

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

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 April 16, 2025

6,935 Δ

jax-v0.7.1 Released August 19, 2025

19,270 Δ

jax-v0.7.0 Released July 22, 2025

57,104 Δ

jax-v0.6.2 Released June 17, 2025

34,451 Δ

jax-v0.6.1 Released June 18, 2025

50,810 Δ

jax-v0.6.0 Released April 22, 2025

58,237 Δ

Top Contributors in jax-v0.7.1

bchetioui

melissawm

allanrenucci

jakevdp

yashk2810

justinjfu

danielsuo

superbobry

dimitar-asenov

emilyfertig

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.