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
1
Commit Activity
Directory Browser
Velocity History
Releasesjax-v0.4.38
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:
17,793 Diff Delta
Contributors:
56 total committers
Data confidence:
Commits:
234 impacting commits

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
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
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 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 Ξ”
jax-v0.4.35 Released October 22, 2024
0 Ξ”
jax-v0.4.34 Released October 2, 2024
29,039 Ξ”
jax-v0.4.33 Released September 16, 2024
0 Ξ”

Top Contributors in jax-v0.4.38

yashk2810
jakevdp
bchetioui
gnecula
hawkinsp
bixia1
wenscarl
superbobry
dougalm
dfm

Directory Browser for jax-v0.4.38

We haven't yet finished calculating and confirming the files and directories changed in this release. Please check back soon.

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.