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

Jax: jax-v0.3.17 Release

Release date:
August 31, 2022
Previous version:
jax-v0.3.16 (released August 11, 2022)
Magnitude:
20,878 Diff Delta
Contributors:
31 total committers
Data confidence:
Commits:
184 impacting commits

184 Commits in this Release

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

roll-forward #11952
Authored August 20, 2022
introduce key-element-type arrays and overhaul the Python PRNG key array type
Authored August 10, 2022
roll-forward #11952, take 2
Authored August 22, 2022
Initial implementation of eager pmap
Authored August 3, 2022
Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.
Authored August 30, 2022
[JAX] Prepare not to export jax._src by default.
Authored August 23, 2022
teach jax2tf about custom eltypes, key arrays, and random key primitives
Authored August 17, 2022
Cleaning up eager pmap implementation
Authored August 11, 2022
delete old remat implementation
Authored August 16, 2022
sketch: setup for new key array implementation based on eltypes
Authored July 28, 2022
Add support for doing grad of `pjit` (similar to what `jit` supports). Resolve in_shardings in `_pjit_call_impl` (that were UNSPECIFIED) before lowering to XLA. Then check if the device assignment is same across shardings in `lower_sharding_computation`.
Authored August 17, 2022
Enable `jax_array` on CPU for the entire JAX test suite!
Authored August 19, 2022
[jax2tf] Improves support for lax.gather (enable_xla=False).
Authored August 26, 2022
[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.
Authored August 25, 2022
Copybara import of the project:
Authored August 18, 2022
Sparse direct solver using QR factorization from cuSOLVER. This is the jaxlib implementation. We will want to combine this with the sparse libraries already existing in JAX.
Authored August 17, 2022
remove custom_jvp_call_jaxpr_p and its rules
Authored January 8, 2022
Make all pmap tests pass with Array! I am skipping all soft pmap tests for now.
Authored August 12, 2022
Add checkpointing support for Array similar to GDA.
Authored August 22, 2022
Make checkify tests pass with Array and add methods on Array that are present on DA.
Authored August 16, 2022
Revert: Sparse direct solver using QR factorization from cuSOLVER. This is the jaxlib implementation. We will want to combine this with the sparse libraries already existing in JAX.
Authored August 18, 2022
Fixes for PocketFFT->ducc migration.
Authored August 26, 2022
Make eager pmap tests pass with `Array`. Also add a slow path for Array in `pmap` similar to what SDA has. This is required for eager pmap. Adding a slow path removes the need for doing sharding checks in api.py because SDA doesn't do those checks and if the sharding does not match with pmap sharding, then it just defaults to the slow path (exactly like SDA).
Authored August 20, 2022
[jax2tf] lax.reduce_window (enable_xla=False): bug fix and improvements.
Authored August 16, 2022
jax.numpy: improve support for boolean inputs
Authored August 12, 2022
[sparse] add support for bcoo equivalent of lax.slice
Authored August 25, 2022
Add support for interoperability via dlpack for Array and also make pickle_tests and lax_numpy_test pass with Array.
Authored August 18, 2022
Remove `device_replica_id_map` from the Sharding interface because the standalone function should be more than enough to use. The major use-case of this is for checkpointing and accessing addressable_shards which accesses the standalone function makes it work.
Authored August 29, 2022
Expose pure callback and enable rank polymorphic callbacks
Authored August 16, 2022
roll-forward deletion of custom_jvp_call_jaxpr_p
Authored August 18, 2022
Upgrade logistic (sigmoid) function into a lax primitive.
Authored August 24, 2022
roll-forward deletion of custom_jvp_call_jaxpr_p
Authored August 18, 2022
Internal change
Authored August 18, 2022
Take the pjit XLA compilation path for `Arrays`. In the test, `astype` happens in a sharded fashion without the round trip to host.
Authored August 18, 2022
Allow debug prints in staged out custom derivative functions
Authored August 13, 2022
Remove the fallback to `lower_xla_callable` that exist for `jit(pmap)` cases when `Array` was enabled and add minimal support to `lower_sharding_computation`.
Authored August 30, 2022
Check if the buffer shape matches the excepted shard shape by Array.
Authored August 29, 2022
Working multihost eager pmap
Authored August 11, 2022
Apply suggestions from code review
Authored August 17, 2022
Initial transpose implementation
Authored August 12, 2022

Browse Other Releases

Latest Pending Unreleased 😎
jax-v0.3.20 Released September 28, 2022
0 Ξ”
jax-v0.3.20-rc Released September 28, 2022
415 Ξ”
jax-v0.3.19 Released September 27, 2022
1,365 Ξ”
jax-v0.3.18 Released September 26, 2022
785 Ξ”
jax-v0.3.18-rc Released September 26, 2022
17,611 Ξ”
jax-v0.3.17 Released August 31, 2022
20,878 Ξ”
jax-v0.3.16 Released August 11, 2022
16,661 Ξ”
jax-v0.3.15 Released July 22, 2022
0 Ξ”
jax-v0.3.15-rc Released July 22, 2022
18,821 Ξ”
jax-v0.3.14 Released June 27, 2022
25,712 Ξ”

Top Contributors in jax-v0.3.17

froystig
yashk2810
hawkinsp
mattjj
sharadmv
a-googler
jakevdp
marcvanzee
denizokt
apaszke

Directory Browser for jax-v0.3.17

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

Release Notes Published

  • GitHub commits.
  • Bugs
    • Fix corner case issue in gradient of lax.pow with an exponent of zero (#12041)
  • Breaking changes
    • jax.checkpoint, also known as jax.remat, no longer supports the concrete option, following the previous version's deprecation; see JEP 11830.
  • Changes
    • Added jax.pure_callback that enables calling back to pure Python functions from compiled functions (e.g. functions decorated with jax.jit or jax.pmap).
  • Deprecations:
    • The deprecated DeviceArray.tile() method has been removed. Use jax.numpy.tile (#11944).
    • DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead.
GitClear uses cookies to ensure you get the best experience on our website. Learn more