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:
May 25, 2025
Previous version:
jax-v0.6.0 (released April 22, 2025)
Magnitude:
50,833 Diff Delta
Contributors:
82 total committers
Data confidence:
Commits:
625 impacting commits

625 Commits in this Release

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

[Pallas] Update TPU pipelining docs
Authored April 17, 2025
Update README.md to remove jax.pmap
Authored April 24, 2025
[Pallas] Generalize BlockSpec to support different indexing mode for each dim in the block shape
Authored April 18, 2025
Remove pspec -> names conversion during `shard_map_p.bind` and instead preserve partition specs everywhere internally.
Authored May 20, 2025
initial commit of boxes/lists
Authored April 3, 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. It's implemented as split dq and dkv kernels to have enough SMEM for double-buffering.
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
[JAX] Make fully replicated sharding to avoid materializing the same host buffers
Authored May 13, 2025
Fix bugs wrt partial-auto when there are multiple levels of nesting.
Authored April 30, 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

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 April 16, 2025
452,158 Ξ”
jax-v0.6.1 Released May 25, 2025
50,833 Ξ”
jax-v0.6.0 Released April 22, 2025
57,866 Ξ”
jax-v0.5.3 Released March 31, 2025
42,365 Ξ”
jax-v0.5.2 Released March 4, 2025
21 Ξ”
jax-v0.5.1 Released April 5, 2025
61,594 Ξ”
jax-v0.5.0 Released February 15, 2025
41,443 Ξ”

Top Contributors in jax-v0.6.1

yashk2810
mattjj
dfm
hawkinsp
apaszke
justinjfu
superbobry
sharadmv
jakevdp
a-googler

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.
GitClear uses cookies to ensure you get the best experience on our website. Learn more