Home » Oreps » Jax » All Releases » jax-v0.3.5 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 Release Notes

Jax jax-v0.3.5

Jax: jax-v0.3.5 Release

Release date:

April 7, 2022

Previous version:

jax-v0.3.4 (released March 18, 2022)

Magnitude:

9,132 Diff Delta

Contributors:

33 total committers

Data confidence:

Commits:

131

131 Commits in this Release

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

[sparse] Change call signature of bcoo primitive wrappers.

Authored March 29, 2022

First attempt to enable auto-sharding. This CL adds support for GDA (no SDA support yet).

Authored April 1, 2022

Checkify: add way to embed runtime info in error messages.

Authored March 8, 2022

Fix for hipsparse in ROCm.

Authored March 17, 2022

add random.loggamma and improve dirichlet & beta implementation

Authored March 21, 2022

[MHLO] Add MHLO lowerings for parallel collectives.

Authored April 7, 2022

Factor-out reductions from lax_numpy.py

Authored March 18, 2022

Checkify: add axis and axis size to OOB error message.

Authored March 30, 2022

Add istft to jax.scipy.signal.

Authored March 28, 2022

track input argument information in one tree at each AOT stage

Authored March 19, 2022

Remove the jax_enable_mlir flag. MLIR is now the only supported code path.

Authored April 4, 2022

Speed up GDA initialization by making local_shards lazy and hiding checks behind `config.jax_enable_checks` flag.

Authored March 29, 2022

allow for recursive uses of `custom_transpose`

Authored March 25, 2022

require consistent output structure in custom vmap rules

Authored January 28, 2022

[jax2tf] Updates `custom_assert` for jax2tf SVD (primitive) limitations.

Authored March 30, 2022

[MHLO] Prefer backend-specific HLO lowerings instead of non-backend-specific MHLO lowerings.

Authored April 6, 2022

Add jax.scipy.linalg.rsf2csf

Authored March 27, 2022

Deprecate jax.tree_util.tree_multimap

Authored April 1, 2022

improve remat transpose caching (cf. #9661)

Authored March 25, 2022

random_test: add tests of random values for distributions

Authored March 23, 2022

protocols for the lowering and executable underlying `Lowered` and `Compiled`

Authored March 22, 2022

introduce an upgrade flag for `custom_vjp` overhaul (still experimental)

Authored March 26, 2022

Fix race condition for weakref destructor by catching rare exceptions.

Authored March 31, 2022

[sparse] consolidate flavors of bcoo_dot_general

Authored April 7, 2022

* Disallow any other type other than GDA and ShapedArray for auto sharding.

Authored April 4, 2022

Migrate 'jaxlib' CPU custom-calls to the status-returning API

Authored March 30, 2022

Separate jax.test_util implementations into public and private sources.

Authored April 4, 2022

Remove sharded_jit usage from jax2tf because sharded_jit is deprecated.

Authored April 7, 2022

[sparse] bcoo_mul: support mixing batch & sparse dims

Authored March 18, 2022

[sparse] add bcoo_add_batchdim

Authored April 6, 2022

[mhlo] Clean up ops that can use InferTensorTypeWithReify

Authored March 30, 2022

Consolidation of hipsolver/cusolver APIs.

Authored April 6, 2022

[shape_poly] Change internal representation of parametric shapes

Authored March 28, 2022

jax2tf: Support uint32 keys in rng_bit_generator.

Authored April 7, 2022

doc: add examples of using partial with jit

Authored March 29, 2022

Add weakreaf_lru_cache to prevent caches from pinning jaxprs.

Authored March 19, 2022

When passing `ShapedArray` as input to pjit.lower() (which is not allowed normally i.e `pjit(f)(*args)`), consider it as Global. This will also help during the auto sharding change.

Authored March 24, 2022

avoid batching units in cond partial eval

Authored March 23, 2022

[MHLO] Add MHLO lowering for PRNG kernels.

Authored April 6, 2022

refine occurrences of the "forward-mode disallowed" error under new `custom_vjp`

Authored March 26, 2022

Browse Other Releases

Latest Pending Unreleased 😎

jax-v0.3.9 Released May 3, 2022

878 Δ

jax-v0.3.8 Released April 30, 2022

11,149 Δ

jax-v0.3.7 Released April 15, 2022

4,073 Δ

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 Δ

Top Contributors in jax-v0.3.5

jakevdp

yashk2810

froystig

hawkinsp

tlu7

LenaMartens

reza-rahimi-144f

yotarok

mattjj

pschuh

Directory Browser for jax-v0.3.5

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

Release Notes Published

Changes * added jax.random.loggamma & improved behavior of jax.random.beta and jax.random.dirichlet for small parameter values (#9906). * the private lax_numpy submodule is no longer exposed in the jax.numpy namespace (#10029). * added array creation routines jax.numpy.frombuffer, jax.numpy.fromfunction, and jax.numpy.fromstring (#10049). * DeviceArray.copy() now returns a DeviceArray rather than a np.ndarray (#10069) * added jax.scipy.linalg.rsf2csf * Deprecations: * jax.nn.normalize is being deprecated. Use jax.nn.standardize instead (#9899). * jax.tree_util.tree_multimap is deprecated. Use jax.tree_util.tree_map instead (#5746). * jax.experimental.sharded_jit is deprecated. Use pjit instead.