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

Jax: jax-v0.3.1 Release

Release date:
February 18, 2022
Previous version:
jax-v0.3.0 (released February 10, 2022)
Magnitude:
6,721 Diff Delta
Contributors:
20 total committers
Data confidence:
Commits:
67 impacting commits

67 Commits in this Release

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

Add some functions for spectral analysis.
Authored February 3, 2022
JaxTestCase: set numpy_rank_promotion='raise' by default
Authored February 1, 2022
Copybara import of the project:
Authored February 11, 2022
Xmap GDA integration. Non-contiguous mesh is allowed!
Authored February 17, 2022
improve pjit in/out_axis_resources pytree errors
Authored February 8, 2022
Try adding support for nesting with_sharding_constraint in MANUAL xmaps
Authored February 15, 2022
[remove units] make JaxprTrace.process_call not introduce units
Authored February 7, 2022
DOC: add ability to document extra_params within _wraps
Authored February 10, 2022
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
Authored February 14, 2022
Merge `mesh` and `Mesh`. Make `Mesh` a context manager + class so that it can be used in the following ways:
Authored February 17, 2022
Implement the JAX transfer guard API
Authored February 14, 2022
Add a partial_eval_jaxpr_custom_rule for xmap
Authored February 15, 2022
* Make _old_env thread local so that it can be used in multiple threads.
Authored February 18, 2022
implement `jnp.expand_dims` and `jnp.stack` for PRNGKeyArrays
Authored February 17, 2022
Turn execute_replicated into a class so we can access its fields.
Authored February 18, 2022
Add multi-host utilities to JAX core. Adapted from https://github.com/google-research/t5x/blob/main/t5x/multihost_utils.py
Authored February 15, 2022
Add translation rule for optimization barrier.
Authored February 14, 2022
[sparse] Updates `bcoo_dot_general` cuSparse lowering rule by adding sorted indices.
Authored February 14, 2022
Add `ndim` and `size` to GDA
Authored February 11, 2022
adds jax.scipy.linalg.sqrtm
Authored February 12, 2022
Add deprecation warning to JaxTestCase and JaxTestLoader
Authored February 17, 2022
fix: use boolean tf functions where possible instead of casting to int8 for jax2tf
Authored February 16, 2022
Clarify the NVidia driver version requirements.
Authored February 17, 2022
[JAX] Fix crash when applying jit() to a callable that is not weak-referenceable.
Authored February 15, 2022
add flag for jaxpr colorful syntax highlighting
Authored February 5, 2022
Checkify: fix check_error of nd-error.
Authored February 15, 2022
Checkify: handle named_call in process_call.
Authored February 14, 2022
tests: use lax.broadcast_shapes in place of custom logic
Authored February 10, 2022
Respect __jax_array__ in jnp.ndarray operations
Authored February 11, 2022
Add constant handler for tokens.
Authored February 14, 2022
Update compilation cache logging to log individual hashes as well as cumulative hash
Authored February 8, 2022
adds jax.scipy.schur
Authored February 12, 2022
Rename wrapper functions to always refer to the JAX api function.
Authored February 8, 2022
Allow unevenly partitioned sharding_constraints.
Authored February 18, 2022
Add jax.distributed and jax.dlpack to the docs.
Authored February 17, 2022
err on repeated axes to `expand_dims`, as numpy does
Authored February 17, 2022
Recommend optax in jax.experimental_libraries.optimizers documentation.
Authored February 7, 2022
Revert back to adding aval on Device buffers inside local_shards and convert the cached property to just the normal property.
Authored February 18, 2022
[jax2tf] Fixed stale documentation about XLA metadata.
Authored February 16, 2022
[jax2tf] Added more links to the documentation
Authored February 18, 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,147 Ξ”
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,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 Ξ”

Top Contributors in jax-v0.3.1

jakevdp
yashk2810
mattjj
yotarok
apaszke
a-googler
hawkinsp
hyeontaek
LenaMartens
froystig

Directory Browser for jax-v0.3.1

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

Release Notes Published

  • Changes:
    • jax.test_util.JaxTestCase and jax.test_util.JaxTestLoader are now deprecated. The suggested replacement is to use parametrized.TestCase directly. For tests that rely on custom asserts such as JaxTestCase.assertAllClose(), the suggested replacement is to use standard numpy testing utilities such as numpy.testing.assert_allclose(), which work directly with JAX arrays (#9620 ).
    • jax.test_util.JaxTestCase now sets jax_numpy_rank_promotion='raise' by default (#9562 ). To recover the previous behavior, use the new jax.test_util.with_config decorator: python @jtu.with_config(jax_numpy_rank_promotion='allow') class MyTestCase(jtu.JaxTestCase): ...
    • Added jax.scipy.linalg.schur, jax.scipy.linalg.sqrtm, jax.scipy.signal.csd, jax.scipy.signal.stft, jax.scipy.signal.welch.
GitClear uses cookies to ensure you get the best experience on our website. Learn more