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

Jax: jax-v0.6.1 Release

Release date:
June 18, 2025
Previous version:
jax-v0.6.0 (released April 22, 2025)
Magnitude:
62,649 Diff Delta
Contributors:
82 total committers
Data confidence:
Commits:
624 impacting commits

624 Commits in this Release

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

Update README.md to remove jax.pmap
Authored April 24, 2025
initial commit of boxes/lists
Authored April 3, 2025
Remove pspec -> names conversion during `shard_map_p.bind` and instead preserve partition specs everywhere internally
Authored May 20, 2025
[Pallas] Generalize BlockSpec to support different indexing mode for each dim in the block shape
Authored April 18, 2025
[hijax] landing prototype pieces
Authored May 8, 2025
[JAX] Use `xla::ifrt::Client::MakeArraysFromHostBufferShards()` in Array creation when possible
Authored April 22, 2025
Remove code to support jaxlib < v0.6
Authored April 17, 2025
Replace safe_zip(*) with zip(*, strict=True) in cases where list output is not necessary
Authored April 30, 2025
Remove shard_map rewrite machinery since we have `vma` in types by default
Authored April 17, 2025
[pallas] `pl.pallas_call` no longer allows `compiler_params=` to be a param->value dict
Authored April 23, 2025
Add explicit sharding to Introduction to parallel programming doc
Authored April 16, 2025
[platform_dependent] Ensure that platform_dependent only lowers for intended platforms
Authored May 8, 2025
[Pallas TPU] Introduce a BoundedSlice block shape type
Authored April 25, 2025
[XLA:Python] [JAX] Move XlaBuilder bindings out of JAX and into XLA:Python
Authored April 25, 2025
Ensure __jax_array__ works properly with JIT disabled
Authored April 25, 2025
[hijax] type-changing boxes with pytree contents
Authored May 15, 2025
[Mosaic GPU] Adding a deterministic backwards pass to the pallas MGPU kernel
Authored April 24, 2025
Add basic support of `unreduced` to sharding-in-types! We cannot lower it right now, but it atleast shows up in types
Authored May 8, 2025
[Pallas TPU] Add better support for BoundedSlice and other BlockDim types to emit_pipeline and fusions
Authored April 30, 2025
[Mosaic:TPU][Relayout] Row shifts for packed types and non-native tilings
Authored April 23, 2025
[Pallas/Mosaic GPU] Allow explicit `smem` aliasing
Authored April 23, 2025
Update tridiagonal solve kernels on GPU to properly use the FFI
Authored May 13, 2025
[Mosaic GPU] Add layout inference and lowering for `scf.WhileOp` and enable tests
Authored May 13, 2025
Move const folding and forwarding into tracing
Authored April 24, 2025
[Pallas][Mosaic GPU] Add support for blackwell mma instruction
Authored May 2, 2025
[Mosaic TPU] Fold reshape (..., M, N, 128) -> (..., M, N * 128) with any tiling and dtype to load
Authored May 13, 2025
Reverts c2ba1790417ca206a4d88b25aef4d5ae510dd717
Authored April 18, 2025
[Mosaic GPU] Add suport for packed TMEM
Authored April 29, 2025
Consolidate initial/final style custom_vjp primitives into one
Authored May 7, 2025
Fix bugs wrt partial-auto when there are multiple levels of nesting
Authored April 30, 2025
[JAX] Make fully replicated sharding to avoid materializing the same host buffers
Authored May 13, 2025
add mc ptr support to tma with overlapped gemm and all reduce examples
Authored May 12, 2025
[Pallas][Mosaic GPU] Expand TMEM support
Authored May 8, 2025
Reverts 49e25c6167806ba90efe8370fb04db3f4966437c
Authored April 23, 2025
[Mosaic GPU] Add a `BroadcastInDim` op in the mosaic mlir dialect
Authored April 28, 2025
[JAX:Sparse] Implement CSR sparse kernel
Authored April 30, 2025
Enable cross-compilation builds of the wheels
Authored April 17, 2025
[JAX] Remove xla_client_test
Authored April 25, 2025
[JAX] Remove jaxlib/xla directory
Authored April 29, 2025
add docs for multi-process run in Kubernete
Authored May 7, 2025

Browse Other Releases

Latest Pending Unreleased 😎
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 Ξ”
jax-v0.5.1 Released April 5, 2025
73,856 Ξ”

Top Contributors in jax-v0.6.1

yashk2810
mattjj
dfm
hawkinsp
apaszke
superbobry
jakevdp
sharadmv
a-googler
dimitar-asenov

Directory Browser for jax-v0.6.1

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.lax.axis_size which returns the size of the mapped axis given its name.

  • Changes

    • Additional checking for the versions of CUDA package dependencies was reenabled, having been accidentally disabled in a previous release.
    • JAX nightly packages are now published to artifact registry. To install these packages, see the JAX installation guide.
    • jax.sharding.PartitionSpec no longer inherits from a tuple.
    • jax.ShapeDtypeStruct is immutable now. Please use .update method to update your ShapeDtypeStruct instead of doing in-place updates.

  • Deprecations

    • jax.custom_derivatives.custom_jvp_call_jaxpr_p is deprecated, and will be removed in JAX v0.7.0.