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

Jax: jax-v0.5.0 Release

Release date:
February 15, 2025
Previous version:
jax-v0.4.38 (released January 16, 2025)
Magnitude:
43,803 Diff Delta
Contributors:
63 total committers
Data confidence:
Commits:
474 impacting commits

474 Commits in this Release

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

Add ensure_arraylike utility for lax.numpy implementations
Authored January 17, 2025
A test for using symbolic shapes for minformer
Authored December 6, 2024
Enable pivoted QR on CPU devices
Authored December 4, 2024
[JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly
Authored January 8, 2025
[ROCm] Implement RNN support
Authored December 2, 2024
[pallas][mosaic kernel export] Add initial support for exporting a dynamic shapes (placeholder bound) kernel out of mosaic, via pallas as both MLIR and jaxpr
Authored December 17, 2024
Make JAX test suite pass (at least most of the time) with multiple threads enabled
Authored January 10, 2025
Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU
Authored January 10, 2025
[Mosaic] NFC: Pull out vreg related functions to util
Authored January 2, 2025
add segment packing
Authored January 9, 2025
[internal] add deprecation test utilities
Authored January 9, 2025
Switch to a new thread-safe utility for catching warnings
Authored December 20, 2024
Adding GPU paged attention kernel
Authored January 5, 2025
#sdy support JAX export tests when Shardy is enabled
Authored October 16, 2024
[Mosaic GPU] Add initial layout mismatch resolution for `splat`/`strided` layouts
Authored January 10, 2025
[Mosaic GPU] Add a lowering for simple `async_load` and `async_store` ops
Authored December 20, 2024
Add a discussion of sharding to the FFI tutorial
Authored January 7, 2025
[sharding_in_types] Functions like einsum, reshape, broadcast_in_dim, broadcasted_iota, convert_element_type and sharding_cast that take out_sharding as an argument in their signature should also allow `PartitionSpec` instead of just `NamedSharding` as an input
Authored January 8, 2025
Move jex.ffi to jax.ffi
Authored December 20, 2024
[better_errors] Add debug_info to DynamicJaxprTrace and JaxprStackFrame
Authored January 8, 2025
ref errors
Authored October 2, 2024
[Mosaic] NFC: Pull out vreg related functions to util
Authored December 19, 2024
[sharding_in_types] Introduce `auto_mode`, `user_mode`, `auto_mode_ctx` and `user_mode_ctx` as **private** APIs to make writing auto/user sharding in types code way easier and noise-free
Authored January 10, 2025
JEP: effort-based versioning (EffVer)
Authored December 18, 2024
Add Pallas Philox implementation
Authored December 17, 2024
[aot] Add support for as_text(debug_info=True)
Authored January 9, 2025
[mosaic_gpu] Added a serialization pass
Authored January 9, 2025
Add JAX_COMPILATION_CACHE_EXPECT_PGLE option
Authored November 14, 2024
[Mosaic:TPU] Vreg-slice-aligned offset changes with scratch retiling
Authored December 4, 2024
[Mosaic TPU] Add relayout-insertion pass and support bitwidth change for i1 vector relayout
Authored December 20, 2024
[Mosaic GPU] Prototype of a warp-specialized pipeline emitter for Mosaic GPU
Authored December 19, 2024
add mutable array ref error checks to cond and custom_vjp
Authored December 18, 2024
[Mosaic GPU][NFC] Simplify and clean up layout inference tests to use `FuncOp`s
Authored January 16, 2025
[jax:custom_partitioning] Support SdyShardingRule with multiple leading
Authored January 14, 2025
[Mosaic:TPU] For trunc, expand supported tilings, offsets and bitwidths
Authored December 9, 2024
Even more linearize fixes
Authored December 18, 2024
[Mosaic:TPU] For trunc, expand supported tilings, offsets and bitwidths
Authored December 19, 2024
Port tests away from setUpClass and setUpModule to setUp alone
Authored January 8, 2025
[Mosaic GPU] Replace the dialect's layout enum with layouts holding the proper
Authored December 18, 2024
#sdy enable pure callbacks and debug prints in JAX
Authored December 9, 2024

Browse Other Releases

Latest Pending Unreleased 😎
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 Ξ”
jax-v0.4.35 Released October 22, 2024
0 Ξ”
jax-v0.4.34 Released October 2, 2024
29,039 Ξ”

Top Contributors in jax-v0.5.0

jakevdp
hawkinsp
a-googler
yashk2810
dfm
apaszke
gnecula
bchetioui
tlongeri
superbobry

Directory Browser for jax-v0.5.0

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

Release Notes Published

As of this release, JAX now uses effort-based versioning. Since this release makes a breaking change to PRNG key semantics that may require users to update their code, we are bumping the "meso" version of JAX to signify this.

  • Breaking changes

    Two key factors motivated this decision: * The Mac x86 build (only) has a number of test failures and crashes. We would prefer to ship no release than a broken release. * Mac x86 hardware is end-of-life and cannot be easily obtained for developers at this point. So it is difficult for us to fix this kind of problem even if we wanted to.

    We are open to readding support for Mac x86 if the community is willing to help support that platform: in particular, we would need the JAX test suite to pass cleanly on Mac x86 before we could ship releases again.

  • Changes:

    • The minimum NumPy version is now 1.25. NumPy 1.25 will remain the minimum supported version until June 2025.
    • The minimum SciPy version is now 1.11. SciPy 1.11 will remain the minimum supported version until June 2025.
    • jax.numpy.einsum now defaults to optimize='auto' rather than optimize='optimal'. This avoids exponentially-scaling trace-time in the case of many arguments (#25214).
    • jax.numpy.linalg.solve no longer supports batched 1D arguments on the right hand side. To recover the previous behavior in these cases, use solve(a, b[..., None]).squeeze(-1).

  • New Features

    • jax.numpy.fft.fftn, jax.numpy.fft.rfftn, jax.numpy.fft.ifftn, and jax.numpy.fft.irfftn now support transforms in more than 3 dimensions, which was previously the limit. See #25606 for more details.
    • Support added for user defined state in the FFI via the new jax.ffi.register_ffi_type_id function.
    • The AOT lowering .as_text() method now supports the debug_info option to include debugging information, e.g., source location, in the output.

  • Deprecations

    • From jax.interpreters.xla, abstractify and pytype_aval_mappings are now deprecated, having been replaced by symbols of the same name in jax.core.
    • jax.scipy.special.lpmn and jax.scipy.special.lpmn_values are deprecated, following their deprecation in SciPy v1.15.0. There are no plans to replace these deprecated functions with new APIs.
    • The jax.extend.ffi submodule was moved to jax.ffi, and the previous import path is deprecated.

  • Deletions

    • jax_enable_memories flag has been deleted and the behavior of that flag is on by default.
    • From jax.lib.xla_client, the previously-deprecated Device and XlaRuntimeError symbols have been removed; instead use jax.Device and jax.errors.JaxRuntimeError respectively.
    • The jax.experimental.array_api module has been removed after being deprecated in JAX v0.4.32. Since that release, jax.numpy supports the array API directly.