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

Jax: jax-v0.4.31-rc Release

Release date:
July 29, 2024
Previous version:
jax-v0.4.30 (released June 18, 2024)
Magnitude:
43,117 Diff Delta
Contributors:
66 total committers
Data confidence:
Commits:
531 impacting commits

531 Commits in this Release

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

Delete `xmap` and the `jax.experimental.maps` module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).
Authored July 24, 2024
[pallas] Added more documentation for grid and BlockSpec.
Authored June 27, 2024
Don't wrap singleton ir.Values with tuples during HLO lowering.
Authored July 1, 2024
[export] Add documentation for debugging and for ensuring compatibility.
Authored June 19, 2024
[ROCm] Bring-up pjrt support
Authored June 12, 2024
[pallas] Move some tests from tpu/pallas_call_test to pallas_test
Authored July 8, 2024
[pallas] Improve some error messages and add API tests.
Authored July 2, 2024
Add support to pipeline emitter for shapes that don't perfectly divide the block shapes
Authored June 18, 2024
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
Authored July 10, 2024
[Pallas] Implement block-invariant sampling.
Authored June 24, 2024
Add concrete layout API to JAX. The API takes `major_to_minor: tuple[int, ...]` and `tiling: tuple[tuple[int, ...], ...]` as the arguments. Allows users to pass layouts to `with_sharding_constraint` to constrain the layout + sharding.
Authored June 27, 2024
PR #21371: [NVIDIA] Add new SDPA API to jax.nn
Authored July 8, 2024
[JAX] Add caching to pjit._infer_params.
Authored June 21, 2024
Added basic support for int2/uint2 dtypes to JAX
Authored July 4, 2024
[pallas] Move tests to the GPU- and TPU- specific directories
Authored July 2, 2024
[Mosaic GPU] Add a warp specialized kernel with a separate TMA warp
Authored June 28, 2024
Update persistent compilation cache doc
Authored July 4, 2024
Internal: remove one_to_one_unop factory
Authored June 24, 2024
[Pallas] Add limited boolean memref support for scalars.
Authored July 16, 2024
Documented a few more Pallas APIs and added them to the API docs
Authored July 19, 2024
Shmallas, a.k.a. allow lowering shard_map + run_state to a pallas_call.
Authored July 23, 2024
Changed the naming of internal config APIs
Authored April 15, 2024
Clean up
Authored July 3, 2024
[Mosaic GPU] Simplify the matmul example
Authored July 25, 2024
[pallas] Document the indexing_mode=Unblocked()
Authored July 7, 2024
Don't wrap singleton ir.Types during HLO lowering.
Authored July 3, 2024
pjit.py cleanups.
Authored June 20, 2024
_make_dispatch_table calls in Pallas GPU lowering are now slightly more compact
Authored June 18, 2024
Add `sharding` to `convert_element_type_p` primitive.
Authored July 9, 2024
Migrated ``pl.BlockSpec`` uses in JAX to the new argument order
Authored July 2, 2024
Squashed commit of the following:
Authored July 2, 2024
[Pallas][Mosaic] Relax dynamic index on 2nd minor dim in load/store.
Authored July 2, 2024
[jax2tf] Deprecate jax2tf with native_serialization=False or enable_xla=False.
Authored July 16, 2024
[Mosaic GPU] Use explicit WGMMA/ALU scheduling in the flash attention kernel
Authored June 20, 2024
Storing the last access time of a cache entry in a separate file
Authored June 17, 2024
Improve compilation cache tests
Authored June 27, 2024
Move FFI helper macros from jaxlib/cpu/lapack_kernels.cc to a jaxlib/ffi_helpers.h.
Authored July 10, 2024
Port Householder Product to XLA's FFI
Authored July 12, 2024
Remove dead code now that xmap is deleted
Authored July 24, 2024
[pallas] Fix the interpreter for block_shape not dividing the overall shape
Authored July 1, 2024

Browse Other Releases

Latest Pending Unreleased 😎
jax-v0.4.33 Released September 16, 2024
0 Ξ”
jax-v0.4.33-rc Released September 16, 2024
43 Ξ”
jax-v0.4.32 Released September 11, 2024
0 Ξ”
jax-v0.4.32-rc Released September 11, 2024
0 Ξ”
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 Ξ”

Top Contributors in jax-v0.4.31-rc

yashk2810
gnecula
apaszke
superbobry
hawkinsp
jakevdp
rajasekharporeddy
sharadmv
dfm
ayaka14732

Directory Browser for jax-v0.4.31-rc

All files are compared to previous version, jax-v0.4.30. 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