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

Jax: jax-v0.2.26 Release

Release date:
December 8, 2021
Previous version:
jax-v0.2.25 (released October 25, 2021)
Magnitude:
29,514 Diff Delta
Contributors:
27 total committers
Data confidence:
Commits:
281 impacting commits

281 Commits in this Release

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

Adds an initial prototype of an alternate JAX compilation path that emits the MLIR MHLO/CHLO dialects instead of classic XLA HLO.
Authored November 11, 2021
Split jax.interpreters.xla up into three pieces:
Authored November 22, 2021
Merge most of the MLIR JIT dispatch logic into the common primitive and JIT computation path.
Authored November 29, 2021
add internal `vmappable` interface (part 1)
Authored October 6, 2021
Use imports relative to the `jax` package consistently, rather than `.`-relative imports.
Authored November 24, 2021
[shape_poly] Refactor shape_poly to be independent of TF.
Authored November 17, 2021
Restructures jax2tf converters eval directory.
Authored December 8, 2021
add skeleton checkify transformation
Authored October 29, 2021
[MLIR] Add more parallelism and sharding support via MLIR.
Authored December 3, 2021
Split slice/update_slice/gather/scatter out of jax._src.lax.lax into jax._src.lax.slicing.
Authored November 24, 2021
Copybara import of the project:
Authored December 1, 2021
Increase minimum jaxlib version to 0.1.74.
Authored November 18, 2021
[sparse] refactor jax.experimental.sparse
Authored December 2, 2021
bcoo_dot_general: support contraction along batch dimensions
Authored November 12, 2021
refactor pmap staging, lowering, and compilation
Authored November 19, 2021
[sparse] spdot_general: implement many more cases
Authored November 18, 2021
Make lax.reduce_window variadic.
Authored November 22, 2021
Add a prototype IREE backend for JAX.
Authored November 15, 2021
Lift contiguous mesh requirement for GDAs. If all inputs are GDAs, the mesh can be non-contiguous.
Authored December 3, 2021
Adds support for batch dimensions for `lax.gather` when enable_xla=False. This allows us to, for instance, convert `jax.vmap(lax.dynamic_slice)`.
Authored November 2, 2021
[JAX] Added `jax.lax.linalg.qdwh`.
Authored October 29, 2021
[MLIR] Fix CPU test failures for MLIR lowering.
Authored November 30, 2021
Consolidate more XLA-lowering logic between jit, pmap, and xmap.
Authored November 30, 2021
[sparse] Improve type safety of cusparse lowerings
Authored November 17, 2021
[sparse] add todense() primitive for use in sparsify transform
Authored October 26, 2021
Initial GlobalShardedDeviceArray implementation.
Authored November 4, 2021
Adding support for a special value for in_axis_resources (`pjit.FROM_GSDA`) when GSDA is an input.
Authored November 19, 2021
[sparse]: add tracer-based implementation of sparsify
Authored November 4, 2021
Fix a number of bugs in MLIR translation rules.
Authored November 17, 2021
[jax2tf] Force TF compilation for code under jax.jit.
Authored September 7, 2021
Add gather/scatter mode support to jax2tf.
Authored November 2, 2021
Split windowed reductions and their gradients into a separate file inside the lax implementation.
Authored November 18, 2021
generalize jaxpr simplification machinery
Authored November 16, 2021
Move serialization/de-serialization of GDA into jax.
Authored December 7, 2021
Initial GSDA-pjit integration.
Authored November 12, 2021
Move dtype canonicalization out of core.AbstractValue subclasses.
Authored December 7, 2021
Internal change
Authored December 6, 2021
[sparse] add autodiff & transpose rules for csr primitives
Authored November 2, 2021
[MLIR] Make jit(pmap(...)) work in the MLIR lowering.
Authored November 30, 2021
Improve logics for numerical comparison
Authored November 18, 2021

Browse Other Releases

Latest Pending Unreleased 😎
jax-v0.3.1 Released February 18, 2022
6,721 Ξ”
jax-v0.3.0 Released February 10, 2022
5,980 Ξ”
jax-v0.2.28 Released February 2, 2022
7,023 Ξ”
jax-v0.2.27 Released January 18, 2022
11,661 Ξ”
jax-v0.2.26 Released December 8, 2021
29,514 Ξ”
jax-v0.2.25 Released October 25, 2021
2,029 Ξ”
jax-v0.2.24 Released October 19, 2021
298 Ξ”
jax-v0.2.23 Released October 19, 2021
0 Ξ”
jax-v0.2.22 Released October 13, 2021
37 Ξ”
jax-v0.2.21 Released September 23, 2021
10,380 Ξ”

Top Contributors in jax-v0.2.26

hawkinsp
jakevdp
yashk2810
mattjj
gnecula
marcvanzee
a-googler
froystig
tlu7
zhangqiaorjc

Directory Browser for jax-v0.2.26

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

Release Notes Published

  • Bug fixes:

  • Out-of-bounds indices to jax.ops.segment_sum will now be handled with FILL_OR_DROP semantics, as documented. This primarily afects the reverse-mode derivative, where gradients corresponding to out-of-bounds indices will now be returned as 0. (#8634).

  • jax2tf will force the converted code to use XLA for the code fragments under jax.jit, e.g., most jax.numpy functions (#7839).

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