Home » Oreps » Jax » All Releases » jax-v0.4.38 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.4.38

Jax: jax-v0.4.38 Release

Release date:

January 16, 2025

Previous version:

jax-v0.4.37 (released December 12, 2024)

Magnitude:

15,869 Diff Delta

Contributors:

56 total committers

Data confidence:

Commits:

234

234 Commits in this Release

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

Support FP8 for dot_product_attention

Authored December 13, 2024

Migrate JAX MLIR Python dialect extensions to nanobind.

Authored December 13, 2024

Put abstract_mesh on every eqn so that we can preserve it during `eval_jaxpr` and `check_jaxpr` roundtrip.

Authored December 17, 2024

[jax:custom_partitioning] Make SdyShardingRule a user facing class.

Authored December 6, 2024

[Mosaic GPU] Add CUPTI profiler alongside events-based implementation

Authored November 20, 2024

[sharding_in_types] Enforce AxisTypes to always exist if `set_mesh` is used.

Authored December 11, 2024

Add lax.composite primitive

Authored November 21, 2024

Add a simple ragged all-to-all Pallas TPU kernel.

Authored December 11, 2024

[Mosaic TPU] Add relayout-insertion pass and support bitwidth change for i1 vector relayout

Authored December 12, 2024

[sharding_in_types] Default axis_types to `Auto` for all axis_names if user does not set any AxisType. Also resolve some TODOs now that we have a way for user to set the mesh.

Authored December 11, 2024

[Mosaic] Pull out layout-related functions to free-standing functions.

Authored December 13, 2024

Add error message when using custom_vmap with reverse-mode AD, and add docstrings.

Authored December 9, 2024

[jax:custom_partitioning] Make SdyShardingRule a user facing class.

Authored December 10, 2024

[Mosaic GPU] Add support for fast WGMMA layout changes after 8- to 16-bit upcast

Authored November 5, 2024

Cleanup: remove uses of no-op raise_to_shaped

Authored December 12, 2024

[mosaic] Migrated the serialization pass from codegen to `pass_boilerplate.h`

Authored December 12, 2024

[export] Expand exporting to work with AbstractMesh.

Authored December 12, 2024

[Mosaic GPU] Add an initial skeleton for a layout inference pass.

Authored December 11, 2024

Allow tracing and lowering (with lowering_platforms specified) to work with an AbstractMesh. Such a computation cannot be compiled.

Authored December 13, 2024

Deduplicate some GPU plugin definition code.

Authored December 10, 2024

Cleanup: toward merging core.concrete_aval & xla.abstractify

Authored December 17, 2024

[Mosaic GPU] Add layout inference and initial lowering for `vector.{load,store}`.

Authored December 17, 2024

[Mosaic GPU] Add end-to-end lowering example for a pointwise kernel using the dialect and layout inference.

Authored December 17, 2024

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

Authored December 17, 2024

[Mosaic GPU] Add layout inference and initial lowering for `vector.{load,store}`.

Authored December 13, 2024

Cleanup: merge core.concrete_aval & xla.abstractify

Authored December 13, 2024

[mosaic] Migrated the serialization pass from codegen to `pass_boilerplate.h`

Authored December 12, 2024

[Pallas] Pallas documentation cleanup

Authored December 4, 2024

[Mosaic:TPU] Retilings with replicated sublanes on source and increasing tile size

Authored December 7, 2024

Add a `freeze` primitive to delimit ref lifetimes for AD.

Authored December 9, 2024

Implement `process_call` and `process_custom_jvp` for LinearizeTrace

Authored December 13, 2024

[shape_poly] Improve reasoning for >= in presence of == constraints.

Authored December 11, 2024

Introduce `lax.ragged_all_to_all` primitive

Authored December 10, 2024

[Mosaic:TPU] Replicating retilings with increasing tile size for (a) replicated 2nd minor or (b) 32-bit single-row

Authored December 7, 2024

Fixes to direct linearize

Authored December 11, 2024

[export] Try out exporting with AbstractMesh

Authored December 12, 2024

[Mosaic GPU] Add WGMMA to the Mosaic GPU MLIR Dialect.

Authored December 11, 2024

Add utility script and env for running the CI scripts under Docker

Authored December 16, 2024

[shape_poly] Remove old non_negative support.

Authored December 13, 2024

internal: dedupe lax broadcasting logic

Authored December 11, 2024

Browse Other Releases

Latest Pending Unreleased 😎

jax-v0.5.3 Released March 19, 2025

42,198 Δ

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

41,443 Δ

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,790 Δ

jax-v0.4.35 Released October 22, 2024

0 Δ

jax-v0.4.34 Released October 2, 2024

26,451 Δ

jax-v0.4.33 Released September 16, 2024

0 Δ

Top Contributors in jax-v0.4.38

yashk2810

jakevdp

bchetioui

gnecula

hawkinsp

bixia1

wenscarl

dougalm

superbobry

dfm

Directory Browser for jax-v0.4.38

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

Loading File Browser...

Release Notes Published

Changes:
- jax.tree.flatten_with_path and jax.tree.map_with_path are added as shortcuts of the corresponding tree_util functions.
Deprecations
- a number of APIs in the internal jax.core namespace have been deprecated. Most were no-ops, were little-used, or can be replaced by APIs of the same name in jax.extend.core; see the documentation for {mod}jax.extend for information on the compatibility guarantees of these semi-public extensions.
- Several previously-deprecated APIs have been removed, including:
- from jax.core: check_eqn, check_type, check_valid_jaxtype, and non_negative_dim.
- from jax.lib.xla_bridge: xla_client and default_backend.
- from jax.lib.xla_client: _xla and bfloat16.
- from jax.numpy: round_.
New Features
- jax.export.export can be used for device-polymorphic export with shardings constructed with {func}jax.sharding.AbstractMesh. See the jax.export documentation.
- Added jax.lax.split. This is a primitive version of jax.numpy.split, added because it yields a more compact transpose during automatic differentiation.