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 19, 2025
Previous version:
jax-v0.5.2 (released March 4, 2025)
Magnitude:
42,198 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
[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
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. 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

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 April 16, 2025
451,688 Ξ”
jax-v0.6.0 Released April 22, 2025
57,866 Ξ”
jax-v0.5.3 Released March 19, 2025
42,198 Ξ”
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
41,443 Ξ”
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
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.
GitClear uses cookies to ensure you get the best experience on our website. Learn more