Home » Oreps » Jax » All Releases » jax-v0.4.29-rc 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.29-rc
Basic Details Issues Addressed Top Contributors Directory Browser
Jax jax-v0.4.29-rc

Jax: jax-v0.4.29-rc Release

Release date:
June 10, 2024
Previous version:
jax-v0.4.28 (released May 9, 2024)
Magnitude:
27,592 Diff Delta
Contributors:
71 total committers
Data confidence:
Commits:
408 impacting commits

408 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 😎
jax-v0.4.31 Released July 29, 2024
0 Ξ”
jax-v0.4.31-rc Released July 29, 2024
43,117 Ξ”
jax-v0.4.30 Released June 18, 2024
0 Ξ”
jax-v0.4.30-rc Released June 18, 2024
10,316 Ξ”
jax-v0.4.29 Released June 10, 2024
0 Ξ”
jax-v0.4.29-rc Released June 10, 2024
27,592 Ξ”
jax-v0.4.28 Released May 9, 2024
0 Ξ”
jax-v0.4.28-rc Released May 9, 2024
0 Ξ”
jax-v0.4.27 Released May 7, 2024
0 Ξ”
jax-v0.4.27-rc Released May 7, 2024
18,638 Ξ”

Top Contributors in jax-v0.4.29-rc

jakevdp
superbobry
yashk2810
a-googler
gnecula
apaszke
froystig
dfm
hawkinsp
bythew3i

Directory Browser for jax-v0.4.29-rc

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

Loading File Browser...
GitClear uses cookies to ensure you get the best experience on our website. Learn more