Home » Oreps » Jax » All Releases » jax-v0.3.2 Release

google/jax

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

Watch

Basic Details Issues Addressed Top Contributors Directory Browser

Jax jax-v0.3.2

Jax: jax-v0.3.2 Release

Release date:

March 16, 2022

Previous version:

jax-v0.3.1 (released February 18, 2022)

Magnitude:

13,751 Diff Delta

Contributors:

24 total committers

Data confidence:

Commits:

149

149 Commits in this Release

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

Update JAX to use new math libraries in ROCm-5.0.

Authored February 15, 2022

Add separate mechanism for threading name stacks to the lowering

Authored October 28, 2021

[sparse] change call signature of coo primitive wrappers

Authored March 7, 2022

re-implement `custom_transpose` without upfront staging.

Authored January 21, 2022

Increase minimum jaxlib version to 0.3.0.

Authored March 4, 2022

Document jax.nn.initializers.

Authored March 7, 2022

Rollback of:

Authored March 2, 2022

remove `_const` from public `jax.lax` module

Authored March 7, 2022

[sparse] track sorted columns for COO GPU lowerings

Authored March 9, 2022

Set in_positional_semantics should be GLOBAL for fully replicated values to avoid recompilation.

Authored March 15, 2022

lax_numpy.py: factor ufuncs into their own private submodule

Authored March 15, 2022

Factor-out pieces of lax_numpy.py

Authored March 1, 2022

Make `mesh_axes` on GDA strict by only allowing PartitionSpecs to be consistent with pjit.

Authored March 7, 2022

remove `_convert_element_type` from public `jax.lax` module

Authored March 10, 2022

Always treat all mesh axes controlled by xmap as MANUAL

Authored February 22, 2022

Refactor Jaxpr pretty-printing to use a `JaxprPpSettings` named tuple

Authored March 9, 2022

Remove jax.ops.index... functions.

Authored February 17, 2022

Add the input avals to Lowered and Compiled.

Authored March 9, 2022

[sparse]: COO: check for sorted rows before cusparse lowering

Authored March 8, 2022

[JAX] Update ann to use XLA based fallback ApproxTopK.

Authored March 15, 2022

Documentation improvements.

Authored March 8, 2022

Add an API compatibility policy.

Authored February 17, 2022

lax_numpy.py: factor out some common utilities

Authored March 10, 2022

Improve TPU v2 and v3 mesh_utils.create_device_mesh logic.

Authored March 8, 2022

cleanup: remove _constant_like in favor of lax._const

Authored March 2, 2022

convert to doctest

Authored March 3, 2022

Support gamma distribution with PRNGKeys other than threefry2x32.

Authored March 10, 2022

Index update operators: add scatter_apply()

Authored February 18, 2022

Use real Precision type for lax.PrecisionType

Authored March 4, 2022

in tests, compare jnp operations on PRNGKeyArrays to the same on jnp arrays

Authored February 17, 2022

Checkify: address initial feedback.

Authored February 17, 2022

Handle jaxpr constants correctly in MLIR lowering of conditional branches.

Authored March 1, 2022

format nicely

Authored March 4, 2022

Add JAX support for pjit on CPU

Authored March 9, 2022

segment_max: fix identity for boolean dtype

Authored March 15, 2022

improvements to autodidax reduce_sum

Authored March 16, 2022

Use sharded shape to compute aliasing.

Authored March 10, 2022

Fix tests and handle cond consts.

Authored February 22, 2022

fixing scan and other control flow

Authored February 22, 2022

Make resharding of GDA work if the shape is larger than what it was serialized with.

Authored February 19, 2022

Browse Other Releases

Latest Pending Unreleased 😎

jax-v0.3.6 Released April 13, 2022

5,775 Δ

jax-v0.3.5 Released April 7, 2022

9,132 Δ

jax-v0.3.4 Released March 18, 2022

331 Δ

jax-v0.3.3 Released March 17, 2022

679 Δ

jax-v0.3.2 Released March 16, 2022

13,751 Δ

jax-v0.3.1 Released February 18, 2022

6,721 Δ

jax-v0.3.0 Released February 10, 2022

5,978 Δ

jax-v0.2.28 Released February 2, 2022

829 Δ

jax-v0.2.27 Released January 18, 2022

6,502 Δ

jax-v0.2.26 Released December 8, 2021

1,528 Δ

Top Contributors in jax-v0.3.2

jakevdp

reza-rahimi-144f

hawkinsp

froystig

yashk2810

sharadmv

a-googler

jblespiau

mattjj

skye

Directory Browser for jax-v0.3.2

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