Home » Oreps » Jax » All Releases » jaxlib-v0.4.29 Release

google/jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Watch

Commit Activity

Directory Browser

Velocity History

Releasesjaxlib-v0.4.29

Basic Details Issues Addressed Top Contributors Directory Browser Release Notes

Jax jaxlib-v0.4.29

Jax: jaxlib-v0.4.29 Release

Release date:

June 10, 2024

Previous version:

jaxlib-v0.4.28 (released May 9, 2024)

Magnitude:

27,453 Diff Delta

Contributors:

70 total committers

Data confidence:

Commits:

401

401 Commits in this Release

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

[Mosaic GPU] Integrate the ExecutionEngine with the jaxlib GPU plugin

Authored May 30, 2024

Removed noop # type: ignore comments

Authored May 17, 2024

[XLA:Mosaic] Fix infer/apply vector layout rule for terminators (scf::yieldOp, scf::conditionOp).

Authored May 31, 2024

Cholesky rank-1 update kernel for JAX.

Authored May 14, 2024

Improve docs for jax.numpy set-like operations

Authored May 13, 2024

Added a very rough sketch of Mosaic GPU lowering for Pallas

Authored May 14, 2024

Reverts c2a3c0bb80434d89053b43648f88eba22b9bf1fa

Authored June 5, 2024

Refactor kernel function assigment

Authored June 7, 2024

Avoid global state in compilation_cache_test

Authored May 29, 2024

Add jax.make_array_from_process_local_data to create a distributed tensor from host data and supporting scaffolding in sharding to be able to figure out dimensions of host data required.

Authored May 16, 2024

Enable a lot of previously disabled tests now that compute offload works on TPU.

Authored May 24, 2024

Better documentation for jax.tree_util

Authored May 21, 2024

Initial commit for `jax.experimental.compute_on` API.

Authored May 17, 2024

[Mosaic GPU] Implement non-transposing matmul for bf16 x s8 matmul

Authored May 23, 2024

Add quantization support for PagedAttention TPU Pallas kernel.

Authored May 17, 2024

Make attrs work with pytrees

Authored May 23, 2024

tests: set configs with jtu.with_config rather than manually

Authored June 5, 2024

rewrite the slab in the units that we (think we) want

Authored April 3, 2024

Reverts da816d34eaad6a1c6536959ccb4bfee4466c037d

Authored June 6, 2024

Remove `jaxpr` and `name` from `Lowered` because `specialize` already has those. This keeps the abstraction boundary clear. Adapt `export` to use `specialize`.

Authored June 6, 2024

Another attempt to land #20445

Authored May 24, 2024

Slightly rearranged NDIndexer.from_indices_shape and added missing tests

Authored May 29, 2024

Add a jaxpr interpreter for propagating memory kinds to output. It only triggers if we detect multiple memory kinds in the jaxpr.

Authored May 10, 2024

Add support to pipeline emitter for shapes that don't perfectly divide the block shapes

Authored June 5, 2024

Better documentation for several jax.numpy functions

Authored May 3, 2024

Added pl.debug_print() -- a new primitive for printing from Pallas kernels

Authored May 22, 2024

Removed mentions of iree from the test suite

Authored May 23, 2024

Bumped mypy to 1.10.0 and ruff to 0.4.4

Authored May 16, 2024

Removed the double re-exporting of Pallas GPU/TPU APIs

Authored June 10, 2024

avoid deprecated host_callback.id_print in tests.

Authored May 22, 2024

[Mosaic GPU] Prepare for writing warp-specialized kernels

Authored May 12, 2024

[Mosaic GPU] Make the profiler warpgroup aware

Authored June 4, 2024

[Mosaic GPU] Support f16 wgmma accumulator type.

Authored June 3, 2024

[Mosaic GPU] Restore the PTX/PTXAS/SASS dump flags

Authored May 31, 2024

[Mosaic GPU] Add a WGSplatLayout that trivially supports reshape and broadcast.

Authored May 17, 2024

Added a slow (but working!) implementation of layer norm in Pallas via Mosaic GPU

Authored May 17, 2024

[export] Move the export implementation to jax._src.export.

Authored June 9, 2024

[export] Remove old deprecated APIs for jax.experimental.export.

Authored June 7, 2024

[Pallas][Mosaic] Add lowering rules for PRNG ops.

Authored May 24, 2024

[export] Add helped methods to create XLACompatibleShardings for in_shardings and out_shardings

Authored June 3, 2024

Browse Other Releases

Latest Pending Unreleased 😎

0.3.21 Released September 30, 2022

1,287 Δ

0.3.20 Released September 28, 2022

0 Δ

jaxlib-v0.4.31 Released July 29, 2024

41,440 Δ

jaxlib-v0.4.30 Released June 18, 2024

10,316 Δ

jaxlib-v0.4.29 Released June 10, 2024

27,453 Δ

jaxlib-v0.4.28 Released May 9, 2024

0 Δ

jaxlib-v0.4.27 Released May 7, 2024

17,951 Δ

jaxlib-v0.4.26 Released April 3, 2024

16,864 Δ

jaxlib-v0.4.25 Released February 24, 2024

17,299 Δ

jaxlib-v0.4.24 Released February 5, 2024

29,870 Δ

Top Contributors in jaxlib-v0.4.29

jakevdp

superbobry

yashk2810

a-googler

gnecula

apaszke

froystig

dfm

hawkinsp

bythew3i

Directory Browser for jaxlib-v0.4.29

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

Loading File Browser...

Release Notes Published

Bug fixes
- Fixed a bug where XLA sharded some concatenation operations incorrectly, which manifested as an incorrect output for cumulative reductions (#21403).
- Fixed a bug where XLA:CPU miscompiled certain matmul fusions (https://github.com/openxla/xla/pull/13301).
- Fixes a compiler crash on GPU (https://github.com/google/jax/issues/21396).
Deprecations
- jax.tree.map(f, None, non-None) now emits a DeprecationWarning, and will raise an error in a future version of jax. None is only a tree-prefix of itself. To preserve the current behavior, you can ask jax.tree.map to treat None as a leaf value by writing: jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None).