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
1
Commit Activity
Directory Browser
Velocity History
Releasesjax-v0.5.3
Basic Details Issues Addressed Top Contributors Directory Browser Release Notes
Jax jax-v0.5.3

Jax: jax-v0.5.3 Release

Release date:
March 31, 2025
Previous version:
jax-v0.5.2 (released March 4, 2025)
Magnitude:
45,270 Diff Delta
Contributors:
71 total committers
Data confidence:
Commits:
421 impacting commits

421 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
Open source TPU-friendly ragged paged attention kernel
Authored March 6, 2025
[Mosaic GPU] Enable the new transform inference pass in the warpgroup lowering
Authored March 18, 2025
[Pallas] Adds a simple dynamic race detector for TPU interpret mode
Authored March 5, 2025
Define lax.ragged_dot_general and express lax.ragged_dot in terms of it
Authored March 10, 2025
[sharding_in_types] Rework the `axis_types` argument in Mesh and AbstractMesh APIs
Authored March 13, 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
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
Use the same name for aliased Vars when pretty-printing Jaxprs
Authored February 24, 2025
Reverts 10f6edeb496a2eec2a09c2c5cecbe4f8f02452ab
Authored March 1, 2025
AutoPGLE: force-disable graphs less
Authored March 5, 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. 0
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

Browse Other Releases

Latest Pending Unreleased 😎
jax-v0.7.1 Released August 19, 2025
19,862 Ξ”
jax-v0.7.0 Released July 22, 2025
65,989 Ξ”
jax-v0.6.2 Released June 17, 2025
39,793 Ξ”
jax-v0.6.1 Released June 18, 2025
62,649 Ξ”
jax-v0.6.0 Released April 22, 2025
67,287 Ξ”
jax-v0.5.3 Released March 31, 2025
45,270 Ξ”
jax-v0.5.2 Released March 4, 2025
21 Ξ”
jax-v0.5.1 Released April 5, 2025
73,856 Ξ”
jax-v0.5.0 Released February 15, 2025
43,803 Ξ”
jax-v0.4.38 Released January 16, 2025
17,793 Ξ”

Top Contributors in jax-v0.5.3

apaszke
yashk2810
sharadmv
bchetioui
wenscarl
dfm
a-googler
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.