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

Jax: jax-v0.5.3 Release

Release date:

March 19, 2025

Previous version:

jax-v0.5.2 (released March 4, 2025)

Magnitude:

40,549 Diff Delta

Contributors:

72 total committers

Data confidence:

Commits:

419

419 Commits in this Release

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

[Pallas] Add experimental (private for now) API for manual fusion into Pallas kernels

Authored March 4, 2025

Add jax.experimental._mini_mpmd

Authored February 13, 2025

Scaled matmul for mxfp8

Authored February 5, 2025

[Pallas] Adds a simple dynamic race detector for TPU interpret mode.

Authored March 5, 2025

[sharding_in_types] Rework the `axis_types` argument in Mesh and AbstractMesh APIs. The changes are:

Authored March 13, 2025

Define lax.ragged_dot_general and express lax.ragged_dot in terms of it.

Authored March 10, 2025

[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout

Authored March 12, 2025

[Mosaic GPU][NFC] Refactor MMA SMEM descriptor creation

Authored March 5, 2025

[Pallas] Add option for async DMAs in the new TPU interpret mode

Authored February 26, 2025

block_scale_config

Authored February 13, 2025

[Pallas:MGPU] Adjust warpgroup lowering to the recent emit_pipeline changes

Authored February 25, 2025

[Mosaic GPU] Introduce an initial transform inference pass.

Authored March 14, 2025

[Mosaic-GPU] [2/3] Add NVSHMEM support to Mosaic-GPU custom call

Authored February 21, 2025

#sdy support JAX export tests when Shardy is enabled.

Authored March 3, 2025

[Pallas/Mosaic GPU] Add lowerings/layout inference for all the necessary conversion ops when using Warpgroup semantics.

Authored March 10, 2025

[better_errors] Cleanup use of DebugInfo.arg_names and result_paths

Authored February 18, 2025

Remove `_parsed_pspec` from everywhere in JAX except for NamedSharding constructor. I'll do that in the next CL since that has a dependency on C++ so needs guards.

Authored February 27, 2025

[Mosaic GPU] Add support for small RHS tile sizes in WGMMA

Authored February 28, 2025

Add and test support for partitioning of batch dimensions in lax.linalg.

Authored February 27, 2025

Support nvfp4

Authored March 5, 2025

add autodiff rules for jax.lax.ragged_all_to_all collective

Authored March 12, 2025

AutoPGLE: force-disable graphs less

Authored March 5, 2025

Use the same name for aliased Vars when pretty-printing Jaxprs.

Authored February 24, 2025

Reverts 10f6edeb496a2eec2a09c2c5cecbe4f8f02452ab

Authored March 1, 2025

[Mosaic GPU] Add support for tiling stores to refs using small tiling

Authored March 7, 2025

`PRNGKeyArray.aval` should have the correct logical sharding. This required refactoring code so that we don't hit recursion errors.

Authored March 2, 2025

[pallas:mosaic_gpu] Added WG lowering rules for TMA primitives and run_scoped_p

Authored February 25, 2025

Better mosaic lowering for dynamic shapes, extend an interpreter into shape_poly dimexpr and lower them alongside the graph if we are in a dynamic export regime.

Authored March 18, 2025

[Mosaic GPU] Infer layouts (transforms) on memrefs that directly feed into the dialect wgmma op.

Authored March 2, 2025

Reimplement custom_vjp.optimize_remat using custom_dce.

Authored February 27, 2025

[Mosaic GPU] Add a specialized layout for loading 4-bit inputs in WGMMA friendly layouts

Authored March 18, 2025

Move (most) jaxlib linalg custom call registration into JAX.

Authored March 10, 2025

Port many uses of contextlib.contextdecorator to explicit context manager classes.

Authored February 25, 2025

Fix logging_test fails on Linux with NVIDIA Driver only.

Authored March 18, 2025

[Mosaic GPU] Check in WIP persistent matmul kernel

Authored March 7, 2025

[Mosaic GPU] Add support for .cta_group::2 MMA with n=512 on Blackwell

Authored March 11, 2025

Reverts 342cb7b99a09180472823a33c7cdad8a8db77875

Authored March 5, 2025

[pallas:mosaic_gpu] Added support for math functions to the WG lowering

Authored March 10, 2025

[Mosaic GPU] Remove support for large tiles on Blackwell

Authored March 5, 2025

add reduction support in copy_smem_to_gmem

Authored March 12, 2025

Browse Other Releases

Latest Pending Unreleased 😎

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 December 17, 2024

133,799 Δ

jax-v0.5.3 Released March 19, 2025

40,549 Δ

jax-v0.5.2 Released March 4, 2025

21 Δ

jax-v0.5.1 Released February 26, 2025

60,641 Δ

jax-v0.5.0 Released February 15, 2025

40,607 Δ

jax-v0.4.38 Released January 16, 2025

15,869 Δ

jax-v0.4.37 Released December 12, 2024

3,574 Δ

Top Contributors in jax-v0.5.3

apaszke

yashk2810

sharadmv

wenscarl

dfm

a-googler

bchetioui

gspschmid

superbobry

jburnim

Directory Browser for jax-v0.5.3

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
- Added a allow_negative_indices option to jax.lax.dynamic_slice, jax.lax.dynamic_update_slice and related functions. The default is true, matching the current behavior. If set to false, JAX does not need to emit code clamping negative indices, which improves code size.
- Added a replace option to jax.random.categorical to enable sampling without replacement.