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

Basic Details Issues Addressed Top Contributors Directory Browser Release Notes

Jax jax-v0.5.1

Jax: jax-v0.5.1 Release

Release date:

February 26, 2025

Previous version:

jax-v0.5.0 (released February 15, 2025)

Magnitude:

60,641 Diff Delta

Contributors:

85 total committers

Data confidence:

Commits:

664

664 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. There aren't a lot of failures in TGP but we can atleast upstream these changes until we work on the failures.

Authored January 28, 2025

Add ensure_arraylike utility for lax.numpy implementations

Authored January 17, 2025

Scaled matmul for mxfp8

Authored February 5, 2025

[sharding_in_types] Add vmap + explicit sharding support. The main changes are:

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

[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout

Authored January 28, 2025

Don't allow users to query `tracer.sharding` even under sharding in types mode.

Authored January 17, 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

[better_errors] Improvements in propagation of debugging info

Authored January 15, 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. Also some more changes as listed below:

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

[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

Adds a minimal but viable implementation of string arrays (with `numpy.dtypes.StringDType`) in JAX. Currently this only supports making of a string array by means of either `jax.numpy.asarray` or `jax.device_put` and reading it back with `jax.device_get`.

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] Ensure debug_info.arg_names is never None.

Authored January 20, 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

[better_errors] Add more debug info test coverage

Authored January 25, 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] Continue adding debug info to Jaxprs (step 4)

Authored February 6, 2025

refactor: move lax_numpy indexing routines to their own submodule

Authored February 12, 2025

[better_errors] Make it explicit that debug_info is not None.

Authored February 14, 2025

Browse Other Releases

Latest Pending Unreleased 😎

jaxlib-v0.1.36 Released November 22, 2019

0 Δ

jaxlib-v0.1.33 Released December 17, 2024

133,799 Δ

jax-v0.5.3 Released March 19, 2025

40,549 Δ

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

40,607 Δ

jax-v0.4.38 Released January 16, 2025

15,869 Δ

jax-v0.4.37 Released December 12, 2024

3,574 Δ

jax-v0.4.36 Released January 27, 2025

71,795 Δ

jax-v0.4.35 Released October 22, 2024

0 Δ

Top Contributors in jax-v0.5.1

yashk2810

jakevdp

gnecula

apaszke

dfm

dimitar-asenov

a-googler

hawkinsp

superbobry

jburnim

Directory Browser for jax-v0.5.1

All files are compared to previous version, jax-v0.5.0. Click here to browse diffs between other versions.

Loading File Browser...

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.