Home » Oreps » Jax » All Releases » jax-v0.8.2 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.8.2

Jax: jax-v0.8.2 Release

Release date:

December 18, 2025

Previous version:

jax-v0.8.1 (released November 18, 2025)

Magnitude:

19,285 Diff Delta

Contributors:

37 total committers

Data confidence:

Commits:

198

198 Commits in this Release

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

Remove dynamic shapes. Dead weight at this point

Authored November 23, 2025

Refactor indexing code

Authored December 16, 2025

[no-thunks] Implement FlatTree, an internal representation of pytrees

Authored December 10, 2025

[Mosaic GPU] Get rid of hints in layout inference

Authored December 8, 2025

[Mosaic GPU][NFC] Make `Equation` a type of `Constraint`

Authored November 24, 2025

Deprecate jax.ad_checkpoint.remat and jax.ad_checkpoint.checkpoint

Authored December 2, 2025

[Mosaic TPU] Support second/third minor transpose when sublane_count is not 8

Authored November 18, 2025

[Mosaic GPU] Remove `{Least,Most}ReplicatedExpression` constructors

Authored December 9, 2025

[Pallas] Skip interpret tests on accelerators

Authored December 4, 2025

[Mosaic GPU] Add support for all kinds of TMA reductions

Authored December 16, 2025

[no-thunks] Deprecate old form of `cond` and remove some tracing layers

Authored December 2, 2025

[indexing] add strategy argument to rewriting_take

Authored December 1, 2025

[Pallas:SC] Adds a tpu dialect op for sort, and uses it to lower jnp.sort on SC vector subcore

Authored November 26, 2025

Update HLO parser for new HLO format

Authored November 21, 2025

[Mosaic GPU][NFC] Replace deprecated usages of `Builder::create<OpType>` with `OpType::create`

Authored November 27, 2025

Remove type parameters. Type checker doesn't like them

Authored December 11, 2025

[Mosaic GPU] Impelment basic support for `memref.expand_shape` in WG semantics

Authored December 2, 2025

[Pallas:SC] Add plsc.sort_key_val to give access to the `mask` and `descending` args

Authored December 8, 2025

[Pallas:MGPU] Fix `plgpu.inline_mgpu` support for scalar results

Authored November 24, 2025

[pxla] Deprecate `jax.interpreters.pxla` symbols

Authored December 15, 2025

[types] improve internal type annotations for indexing routines

Authored November 25, 2025

Fixes for Jetson Thor devices

Authored December 3, 2025

[Mosaic GPU][NFC] Create TMEM refs using `undefs`

Authored November 19, 2025

[Pallas/Mosaic TPU] Allow non-leading and non-matching batch dimensions in `dot_general`

Authored November 17, 2025

[Mosaic GPU] Support inference of additional layouts for `BroadcastInDimOp`

Authored November 24, 2025

[maint] clean up jnp.arange implementation

Authored December 9, 2025

[Pallas] Add missing docstrings and clean up mis-renderings

Authored December 1, 2025

Adds a `plsc.cummax` primitive to compute the cumulative maximum along the (masked) elements of a vector

Authored November 21, 2025

[no-thunks] Avoid thunking/linear_util in checkify

Authored December 5, 2025

[Pallas:SC] Add support for gathering among vector lanes

Authored December 1, 2025

Integrate Triton up to 8d445186

Authored December 11, 2025

Re-enable TPU tests now that TPU thread stack sizes have been increased

Authored December 5, 2025

[Mosaic GPU][NFC] Consolidate duplicated logic when checking that all layouts are assigned

Authored November 26, 2025

Add jax_check_static_indices configuration

Authored December 4, 2025

[Mosaic] NFC: Move ops in tpu.td to tpu_ops.td

Authored December 12, 2025

[Mosaic GPU][NFC] Remove now unnecessary `hasattr` checks

Authored November 19, 2025

Fix eigh gradient tests by applying phase-normalization to the eigenvectors

Authored November 21, 2025

[Mosaic GPU] Add support for reductions when writing to GMEM via TMA in WG semantics

Authored December 3, 2025

Use ifrt::AttributeMap::Get instead of directly accessing map

Authored December 9, 2025

[Pallas MGPU] Adding support for squeezed block dims in the pipeline BlockSpecs

Authored December 17, 2025

Browse Other Releases

Latest Pending Unreleased 😎

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

438,752 Δ

jax-v0.9.0 Released January 20, 2026

0 Δ

jax-v0.8.3 Released January 28, 2026

28 Δ

jax-v0.8.2 Released December 18, 2025

19,285 Δ

jax-v0.8.1 Released November 18, 2025

12,749 Δ

jax-v0.8.0 Released October 15, 2025

17,602 Δ

jax-v0.7.2 Released September 15, 2025

12,272 Δ

jax-v0.7.1 Released August 19, 2025

18,096 Δ

Top Contributors in jax-v0.8.2

dougalm

jakevdp

allanrenucci

bchetioui

justinjfu

dimitar-asenov

brianwa84

yueshengys

danielsuo

krishnaharidasan

Directory Browser for jax-v0.8.2

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

Loading File Browser...

Release Notes Published

Deprecations
- jax.lax.pvary has been deprecated. Please use jax.lax.pcast(..., to='varying') as the replacement.
- Complex arguments passed to jax.numpy.arange now result in a deprecation warning, because the output is poorly-defined.
- From jax.core a number of symbols are newly deprecated including: call_impl, get_aval, mapped_aval, subjaxprs, set_current_trace, take_current_trace, traverse_jaxpr_params, unmapped_aval, AbstractToken, and TraceTag.
- All symbols in jax.interpreters.pxla are deprecated. These are primarily JAX internal APIs, and users should not rely on them.
Changes:
- jax's Tracer no longer inherits from jax.Array at runtime. However, jax.Array now uses a custom metaclass such isinstance(x, Array) is true if an object x represents a traced Array. Only some Tracers represent Arrays, so it is not correct for Tracer to inherit from Array.
For the moment, during Python type checking, we continue to declare Tracer as a subclass of Array, however we expect to remove this in a future release. * jax.experimental.si_vjp has been deleted. jax.vjp subsumes it's functionality.