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

Jax: jax-v0.5.1 Release

Release date:
April 5, 2025
Previous version:
jax-v0.5.0 (released February 15, 2025)
Magnitude:
73,856 Diff Delta
Contributors:
86 total committers
Data confidence:
Commits:
667 impacting commits

667 Commits in this Release

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

Now that sharding_in_types config flag is True, remove the config and all the conditionals
Authored February 19, 2025
Some lax.linalg housekeeping
Authored February 10, 2025
Start a new TPU interpret mode for Pallas
Authored November 22, 2024
[sharding_in_types] Upstream changes from defaulting sharding_in_types config to True experiment
Authored January 28, 2025
Add ensure_arraylike utility for lax.numpy implementations
Authored January 17, 2025
Scaled matmul for mxfp8
Authored February 5, 2025
[better_errors] Continue adding debug info to Jaxprs (step 3)
Authored February 5, 2025
[sharding_in_types] Add vmap + explicit sharding support. 0
Authored January 29, 2025
[Mosaic GPU] Implement lowerings for `Tile` and `Transpose` transforms from the MLIR dialect
Authored February 17, 2025
[better_errors] Expand debug info testing with eager mode, and MLIR module checking
Authored January 20, 2025
Add skeleton for a multi-pass source mapper for Jaxprs/HLO to jax.experimental
Authored January 29, 2025
Don't allow users to query `tracer.sharding` even under sharding in types mode
Authored January 17, 2025
[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout
Authored January 28, 2025
Enable pivoted QR on GPU via MAGMA
Authored January 10, 2025
[Pallas] Refactor Pallas HLO interpret mode to a standalone file
Authored January 28, 2025
Adding GPU paged attention kernel
Authored January 5, 2025
block_scale_config
Authored February 13, 2025
[sharding_in_types] Make the typing checks and sharding rule checks a little bit less strict when the current or aval mesh is empty/unset
Authored February 12, 2025
[xla:cpu] FFI: Implement Jax callbacks for CPU/GPU devices using XLA FFI
Authored January 8, 2025
[Mosaic] Several fixes/improvements for the new TPU interpret mode
Authored February 15, 2025
Port many uses of contextlib.contextdecorator to explicit context manager classes
Authored January 26, 2025
[better_errors] Improvements in propagation of debugging info
Authored January 15, 2025
[JAX] Implement an initial object API for colocated Python
Authored February 21, 2025
JEP: pytree preservation of dict key order
Authored January 23, 2025
Use ensure_arraylike utility in jax.numpy sources
Authored January 17, 2025
[better_errors] Refactor more uses of partial_eval.tracing_debug_info (part 1)
Authored January 25, 2025
[Mosaic GPU] Add specialized support for some int4 -> bfloat16 casts
Authored January 22, 2025
refactor: import numpy objects directly in jax.numpy
Authored February 14, 2025
Adding sharding support to dynamic masks
Authored February 7, 2025
Fix debug_nans regressions
Authored December 5, 2024
[better_errors] Ensure debug_info.arg_names is never None
Authored January 20, 2025
Adds a minimal but viable implementation of string arrays (with `numpy.dtypes.StringDType`) in JAX
Authored January 16, 2025
[better_errors] Cleanup use of DebugInfo.arg_names and result_paths
Authored February 18, 2025
[Pallas/Mosaic GPU] Add initial support for warpgroup semantics in lowering
Authored February 12, 2025
[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo
Authored January 31, 2025
refactor: move scalar type defs out of lax_numpy.py
Authored February 6, 2025
[sharding_in_types] Add support for svd_p
Authored January 28, 2025
[pallas:mosaic_gpu] Added support for binary/comparison ops with WG semantics
Authored February 20, 2025
[better_errors] Add more debug info test coverage
Authored January 25, 2025
[better_errors] Continue adding debug info to Jaxprs (step 4)
Authored February 6, 2025

Browse Other Releases

Latest Pending Unreleased 😎
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 Ξ”
jax-v0.4.37 Released December 12, 2024
4,817 Ξ”
jax-v0.4.36 Released January 27, 2025
85,862 Ξ”

Top Contributors in jax-v0.5.1

yashk2810
gnecula
jakevdp
apaszke
dfm
dimitar-asenov
a-googler
hawkinsp
superbobry
jburnim

Directory Browser for jax-v0.5.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 an experimental jax.experimental.custom_dce.custom_dce decorator to support customizing the behavior of opaque functions under JAX-level dead code elimination (DCE). See #25956 for more details.
    • Added low-level reduction APIs in {mod}jax.lax: jax.lax.reduce_sum, jax.lax.reduce_prod, jax.lax.reduce_max, jax.lax.reduce_min, jax.lax.reduce_and, jax.lax.reduce_or, and jax.lax.reduce_xor.
    • jax.lax.linalg.qr, and jax.scipy.linalg.qr, now support column-pivoting on CPU and GPU. See #20282 and #25955 for more details.

  • Changes

    • JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES now work as env vars. Before they could only be specified via jax.config or flags.
    • JAX_CPU_COLLECTIVES_IMPLEMENTATION now defaults to 'gloo', meaning multi-process CPU communication works out-of-the-box.
    • The jax[tpu] TPU extra no longer depends on the libtpu-nightly package. This package may safely be removed if it is present on your machine; JAX now uses libtpu instead.

  • Deprecations

    • The internal function linear_util.wrap_init and the constructor core.Jaxpr now must take a non-empty core.DebugInfo kwarg. For a limited time, a DeprecationWarning is printed if jax.extend.linear_util.wrap_init is used without debugging info. A downstream effect of this several other internal functions need debug info. This change does not affect public APIs. See https://github.com/jax-ml/jax/issues/26480 for more detail.

  • Bug fixes

    • TPU runtime startup and shutdown time should be significantly improved on TPU v5e and newer (from around 17s to around 8s). If not already set, you may need to enable transparent hugepages in your VM image (sudo sh -c 'echo always > /sys/kernel/mm/transparent_hugepage/enabled'). We hope to improve this further in future releases.
    • Persistent compilation cache no longer writes access time file if JAX_COMPILATION_CACHE_MAX_SIZE is unset or set to -1, i.e. if the LRU eviction policy isn't enabled. This should improve performance when using the cache with large-scale network storage.