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

Jax: jax-v0.8.1 Release

Release date:
November 18, 2025
Previous version:
jax-v0.8.0 (released October 15, 2025)
Magnitude:
16,240 Diff Delta
Contributors:
38 total committers
Data confidence:
Commits:
222 impacting commits

222 Commits in this Release

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

[Mosaic GPU][NFC] Rename `OperandOrResult` to `ValueSite`
Authored October 29, 2025
[Pallas:MGPU] Add Pallas lowering for TMEM slices under WG semantic
Authored November 13, 2025
Add test of documentation coverage
Authored November 17, 2025
[pmap] Deprecate `jax.sharding.PmapSharding`, `jax.device_put_replicated` and `jax.device_put_sharded`
Authored November 14, 2025
Add hooks for the experimental JAX repro module
Authored October 28, 2025
[Mosaic TPU] Allow small sublane tilings in second/third minor transpose for unpacked dtype instead of forcing native tiling, which can be more efficient when the second minor dimension of input tensor is small
Authored October 31, 2025
[Mosaic TPU] Allow arbitrary tilings (including large 2nd minor) in second/third minor transpose for packed dtypes instead of forcing single-sublane tiling, which can avoid unnecessary relayouts in many cases
Authored November 10, 2025
[Mosaic GPU] Refactor `Distinct` constraint in equation system
Authored October 28, 2025
[Mosaic GPU][NFC] Delete the `transform_inference` pass
Authored October 16, 2025
Add instructions for updating jax wheels content
Authored November 6, 2025
[Mosaic] Allow padding in small tiling row shuffle reshape
Authored November 17, 2025
[Pallas:SC] Adds support for `VMEM_SHARED` and `HBM` in `plsc.kernel`/`pl.core_map` (with Scalar and VectorSubcore meshes) via `pl.get_global`
Authored October 29, 2025
[Mosaic GPU] Replace inline assembly with NVVM dialect intrinsics in `tcgen05.py`
Authored October 28, 2025
[Pallas/Mosaic GPU] Add support for `f8` types for WGMMA with lhs in registers
Authored November 11, 2025
Add new atomic `live_devices` API
Authored October 23, 2025
[Pallas:MGPU][NFC] Use `utils.is_signed` to determine signedness
Authored November 12, 2025
[Mosaic GPU][NFC] Consolidate `is_signed` derivation in dialect lowering
Authored November 3, 2025
[Mosaic GPU] Add support for integer WGMMA with LHS in registers under WG semantic
Authored November 11, 2025
[Pallas:MGPU] Add initial warpgroup lowering for `tcgen05_mma`
Authored November 4, 2025
lax.top_k: add axis support
Authored November 5, 2025
[Pallas:MGPU] Add Pallas lowering for `plgpu.broadcasted_iota` under WG semantic
Authored November 12, 2025
Adds an 'add' option to dma_start, which adds to rather than overwriting the destination ref
Authored October 20, 2025
[Pallas:MGPU] Add initial support for TMEM under WG semantic
Authored November 3, 2025
feat(gh-13291): Add Pareto distribution functions: `cdf`, `logcdf`, `logsf`, `sf`, `ppf`, and corresponding tests
Authored February 15, 2025
[Mosaic GPU] Add a `fuel` argument to `infer_layout`
Authored November 4, 2025
[Mosaic GPU][NFC] Add the `Argument` VariableType
Authored October 29, 2025
[Pallas] Skip TPU in pipelining scheduler tests
Authored November 14, 2025
add grad hijax coverage tests
Authored October 30, 2025
Handle overlapping specs in scaled matmul custom partitioning
Authored August 27, 2025
Rename dtypes.bit_width to dtypes.itemsize_bits
Authored November 14, 2025
[Mosaic GPU] Replace `jaxlib.version` guards with `hasattr` guards
Authored November 11, 2025
[Pallas:MGPU] Add support for integer WGMMA under WG semantic
Authored November 6, 2025
Repurpose `tpu_sparsecore` annotation for offloading purposes
Authored October 31, 2025
[Mosaic GPU] Implement `dynamic_gcd` which allows finding the `gcd` of an `int` and an integer `ir.Value`
Authored October 21, 2025
Handle pl.BoundedSlice in reshape pull rule
Authored November 6, 2025
jnp.cov & jnp.corrcoef: add support for the dtype argument
Authored October 15, 2025
[Mosaic TPU] Relax restrictions on packed second/third minor transpose to support 3d+ inputs and arbitrary second/third minor dims
Authored October 16, 2025
Add building tests targets on Windows platform in JAX OSS presubmit
Authored November 11, 2025
Add the separate env var to select a custom user mode driver version to match it with KMD version installed on the execution machine
Authored November 10, 2025
[Mosaic GPU] Support `barrier.arrive` with `orders_tensor_core` in WG semantics
Authored November 14, 2025

Browse Other Releases

Latest Pending Unreleased 😎
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 Ξ”
jax-v0.7.0 Released July 22, 2025
65,989 Ξ”

Top Contributors in jax-v0.8.1

allanrenucci
jakevdp
dimitar-asenov
bchetioui
yueshengys
justinjfu
ybaturina
danielsuo
brianwa84
WindQAQ

Directory Browser for jax-v0.8.1

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

Loading File Browser...

Release Notes Published

  • New features:

    • jax.jit now supports the decorator factory pattern; i.e instead of writing python @functools.partial(jax.jit, static_argnames=['n']) def f(x, n): ... you may write python @jax.jit(static_argnames=['n']) def f(x, n): ...

  • Changes:

    • jax.lax.linalg.eigh now accepts an implementation argument to select between QR (CPU/GPU), Jacobi (GPU/TPU), and QDWH (TPU) implementations. The EighImplementation enum is publicly exported from jax.lax.linalg.
    • jax.lax.linalg.svd now implements an algorithm that uses the polar decomposition on CUDA GPUs. This is also an alias for the existing algorithm on TPUs.

  • Bug fixes:

    • Fixed a bug introduced in JAX 0.7.2 where eigh failed for large matrices on GPU (#33062).

  • Deprecations:

    • jax.sharding.PmapSharding is now deprecated. Please use jax.NamedSharding instead.
    • jx.device_put_replicated is now deprecated. Please use jax.device_put with the appropriate sharding instead.
    • jax.device_put_sharded is now deprecated. Please use jax.device_put with the appropriate sharding instead.
    • Default axis_types of jax.make_mesh will change in JAX v0.9.0 to return jax.sharding.AxisType.Explicit. Leaving axis_types unspecified will raise a DeprecationWarning.
    • jax.cloud_tpu_init and its contents were deprecated. There is no reason for a user to import or use the contents of this module; JAX handles this for you automatically if needed.