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:

8,706 Diff Delta

Contributors:

33 total committers

Data confidence:

Commits:

132

132 Commits in this Release

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

[MHLO] Add direct MHLO lowerings for most linear algebra kernels

Authored April 6, 2022

[sparse] Change call signature of bcoo primitive wrappers

Authored March 29, 2022

First attempt to enable auto-sharding. 0

Authored April 1, 2022

Fix for hipsparse in ROCm

Authored March 17, 2022

Checkify: add way to embed runtime info in error messages

Authored March 8, 2022

[MHLO] Add MHLO lowerings for parallel collectives

Authored April 7, 2022

add random.loggamma and improve dirichlet & beta implementation

Authored March 21, 2022

Factor-out reductions from lax_numpy.py

Authored March 18, 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

Checkify: add axis and axis size to OOB error message

Authored March 30, 2022

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

Authored March 29, 2022

Remove the jax_enable_mlir flag. 0

Authored April 4, 2022

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

Authored April 6, 2022

allow for recursive uses of `custom_transpose`

Authored March 25, 2022

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

Authored March 30, 2022

require consistent output structure in custom vmap rules

Authored January 28, 2022

Add jax.scipy.linalg.rsf2csf

Authored March 27, 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

Deprecate jax.tree_util.tree_multimap

Authored April 1, 2022

[sparse] consolidate flavors of bcoo_dot_general

Authored April 7, 2022

improve remat transpose caching (cf. #9661)

Authored March 25, 2022

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

Authored March 30, 2022

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

Authored April 4, 2022

[mhlo] Clean up ops that can use InferTensorTypeWithReify

Authored March 30, 2022

Consolidation of hipsolver/cusolver APIs

Authored April 6, 2022

Add weakreaf_lru_cache to prevent caches from pinning jaxprs

Authored March 19, 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

[sparse] add bcoo_add_batchdim

Authored April 6, 2022

[sparse] bcoo_mul: support mixing batch & sparse dims

Authored March 18, 2022

Remove sharded_jit usage from jax2tf because sharded_jit is deprecated

Authored April 7, 2022

random_test: add tests of random values for distributions

Authored March 23, 2022

Expose AutoSharding's mesh_shape and mesh_ids options to JAX

Authored April 1, 2022

Implement `jnp.trace` in terms of `jnp.diagonal`

Authored April 6, 2022

Revert: implement `jnp.trace` in terms of `jnp.diagonal`

Authored April 6, 2022

doc: add examples of using partial with jit

Authored March 29, 2022

avoid batching units in cond partial eval

Authored March 23, 2022

Avoid retracing when a host_callback.call is called multiple times with the same function

Authored April 1, 2022

Browse Other Releases

Latest Pending Unreleased 😎

jax-v0.3.9 Released May 3, 2022

0 Δ

jax-v0.3.8 Released April 30, 2022

5 Δ

jax-v0.3.7 Released April 15, 2022

358 Δ

jax-v0.3.6 Released April 13, 2022

5,312 Δ

jax-v0.3.5 Released April 7, 2022

8,706 Δ

jax-v0.3.4 Released March 18, 2022

271 Δ

jax-v0.3.3 Released March 17, 2022

609 Δ

jax-v0.3.2 Released March 16, 2022

13,073 Δ

jax-v0.3.1 Released February 18, 2022

6,870 Δ

jax-v0.3.0 Released February 10, 2022

5,168 Δ

Top Contributors in jax-v0.3.5

jakevdp

hawkinsp

yashk2810

froystig

tlu7

LenaMartens

reza-amd

yotarok

mattjj

a-googler

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.