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
1
Commit Activity
Directory Browser
Velocity History
Releasesjax-v0.8.2
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:
22,685 Diff Delta
Contributors:
39 total committers
Data confidence:
Commits:
202 impacting commits

202 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
Add Pallas core_map guide
Authored December 3, 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
[pxla] Deprecate `jax.interpreters.pxla` symbols
Authored December 15, 2025
[types] improve internal type annotations for indexing routines
Authored November 25, 2025
[Pallas:MGPU] Fix `plgpu.inline_mgpu` support for scalar results
Authored November 24, 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

Browse Other Releases

Latest Pending Unreleased 😎
jax-v0.9.2 Released March 18, 2026
15,566 Ξ”
jax-v0.9.1 Released February 28, 2026
19,549 Ξ”
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 Ξ”
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 Ξ”

Top Contributors in jax-v0.8.2

dougalm
jakevdp
allanrenucci
bchetioui
justinjfu
dimitar-asenov
brianwa84
IvyZX
yueshengys
danielsuo

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.