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

Jax: jax-v0.6.2 Release

Release date:
June 17, 2025
Previous version:
jax-v0.6.1 (released June 18, 2025)
Magnitude:
39,793 Diff Delta
Contributors:
67 total committers
Data confidence:
Commits:
428 impacting commits

428 Commits in this Release

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

Add a general system for keeping track of quasi-dynamic data (QDD)
Authored May 30, 2025
Add experimental array serialization for nested pytrees
Authored May 30, 2025
[hijax] all pre-existing Box tests passing, still using typechange env
Authored May 17, 2025
[Pallas] Simulate multiple cores per device when interpreting kernels on CPU
Authored May 22, 2025
[imports] avoid top-level imports in jax.numpy sources
Authored June 3, 2025
[Mosaic GPU] Convert all memrefs with transforms to unrealized casts and check them
Authored June 13, 2025
[Pallas/Mosaic GPU] Add a Pallas:MGPU implementation of `ragged_dot`
Authored May 22, 2025
Clean up some unused GPU linear algebra kernels
Authored June 2, 2025
Add initial support aligned jnp.swapaxes on major/minor dims
Authored May 22, 2025
Move NamedSharding.__eq__ and NamedSharding.__hash__ into C++
Authored June 12, 2025
[Mosaic GPU] Add non-collective blackwell matmul example
Authored May 29, 2025
[pallas] In TPU interpret mode, run kernels in parallel over Megacore cores
Authored June 3, 2025
Refactor jax/_src/api.py and associated files in preparation for moving them to their own BUILD rule
Authored May 30, 2025
[Pallas Fuser] Add support for basic PRNG op fusion
Authored May 29, 2025
[Pallas:MGPU] Add a first prototype of an all_gather collective matmul kernel
Authored May 22, 2025
Fix typos discovered by codespell
Authored May 31, 2025
Initial commit to make unreduced + AD work
Authored June 10, 2025
Port PartitionSpec to C++
Authored June 6, 2025
rename layout to format, part 1
Authored May 29, 2025
PR #28102: Add cudnn paged attention support in JAX cuDNN SDPA API
Authored June 16, 2025
[Mosaic GPU] Resolve different tile transforms using the largest common divisor
Authored June 13, 2025
Move jax/_src/export to its own build rule
Authored June 11, 2025
[Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM)
Authored May 7, 2025
Add is_leaf_with_path predicate
Authored June 10, 2025
[XProf] Change tensorboard-plugin-profile to new xprof package
Authored May 30, 2025
[Pallas] Require parallel dimensions to form a prefix of the grid in TPU interpret mode
Authored May 27, 2025
[Pallas TPU][NFC] Use register to track buffer slots in pipeline loop
Authored June 10, 2025
[Mosaic GPU] Add tests for the Blackwell matmul kernel
Authored May 27, 2025
Reverts b7833e94c1940ed475dae1f5e83e2a984cda5cea
Authored June 9, 2025
[pallas] Slightly revamped how checkify is exposed in Pallas
Authored May 23, 2025
[JAX] Add `vma` to `ShapeDtypeStruct` constructor arguments
Authored June 10, 2025
Add `all_gather_invariant` to lax
Authored June 12, 2025
Add support for non-power-of-2 head size in flash attention
Authored May 22, 2025
[Pallas:MGPU] Make semaphores compatible with plgpu.kernel and the profiler
Authored May 22, 2025
extend pallas paged_attention with kv scales
Authored June 9, 2025
[Mosaic GPU] Add support for async copies to peer devices
Authored May 23, 2025
[Pallas Fuser] Add basic reshape push rule
Authored June 5, 2025
[ragged-paged-attn] Implement static kv cache quantization. 0
Authored May 23, 2025
[Pallas TPU] Support memory space constraints on pallas_call inputs
Authored June 12, 2025
[Pallas][Mosaic GPU] Add collective (CTA-pair) MMAs to blackwell matmul kernel
Authored June 7, 2025

Browse Other Releases

Latest Pending Unreleased 😎
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 Ξ”
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 Ξ”

Top Contributors in jax-v0.6.2

jakevdp
a-googler
apaszke
hawkinsp
yashk2810
justinjfu
mattjj
rdyro
dimitar-asenov
superbobry

Directory Browser for jax-v0.6.2

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 jax.tree.broadcast which implements a pytree prefix broadcasting helper.

  • Changes

    • The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12.