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

Jax: jax-v0.3.16 Release

Release date:
August 11, 2022
Previous version:
jax-v0.3.15 (released July 22, 2022)
Magnitude:
16,659 Diff Delta
Contributors:
29 total committers
Data confidence:
Commits:
167 impacting commits

167 Commits in this Release

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

Add debugging tools under `jax.debug` and documentation
Authored July 26, 2022
prototype unfettered element types in jaxpr arrays
Authored August 5, 2022
Delete sharded_jit
Authored July 29, 2022
Enable receives in TPU callbacks and add tests
Authored August 8, 2022
Refactor state out of `for_loop`
Authored August 1, 2022
[jax2tf] Adds more support for reduce_window when enable_xla=False.
Authored August 11, 2022
fix ad_checkpoint.checkpoint caching issue
Authored July 29, 2022
Remove dependencies on masking.py
Authored July 24, 2022
Make the `_check_shapes_against_resources` check general for all `XLACompatibleSharding`s by looking at the opsharding proto of the shardings.
Authored July 25, 2022
remove control example
Authored July 26, 2022
add cond dce rule and custom-policy partial eval rule
Authored July 19, 2022
tweak checkify docs.
Authored July 29, 2022
Add help docstrings for debugger
Authored July 26, 2022
Add `OpShardingSharding` and add a function that can calculate indices from an opsharding proto
Authored July 29, 2022
Pmap should output SDA like `Array`s to maintain the current behavior exactly. Split the shard_arg_handler for `Array` based on whether the mode is pmap or pjit. Why do this? The doc below explains more about the context.
Authored August 11, 2022
Document multiprocess GPU support.
Authored August 8, 2022
Make LOBPCG test plots compatible with bazel.
Authored August 5, 2022
Create a `SameDeviceAssignmentTuple` type to cache the op shardings and device assignment. But the device_assignment is only cached once because `pjit` checks if all device assignments are equal or not right at the start.
Authored August 11, 2022
Increase minimum NumPy version to 1.20.
Authored August 6, 2022
set new jax.remat / jax.checkpoint to be on-by-default
Authored August 10, 2022
Use XLA extension tokens instead of output tokens
Authored August 4, 2022
Add early support in pjit for single device shardings. Also lift the restriction of needing the mesh context manager when `config.jax_array` is enabled.
Authored August 6, 2022
[jax2tf] Improved documentation and tests for pjit
Authored July 26, 2022
Annotate nn.initializers
Authored February 16, 2022
Annotate nn.initializers
Authored August 4, 2022
small tweaks for bint ad
Authored August 5, 2022
add custom-policy partial eval and dce rules for pmap
Authored July 29, 2022
Enable debug callbacks in checkpoint
Authored August 5, 2022
Don't depend on mesh for `UNSPECIFIED`. Use `OpShardingSharding` for that since its now available and pjit accepts it.
Authored August 5, 2022
Various debugger improvements
Authored August 8, 2022
Fix ConcretizationError in nested calls.
Authored July 6, 2022
Add webpdb option
Authored July 12, 2022
add while_loop custom-policy partial eval rule
Authored July 29, 2022
Change jaxdb->jdb and add option to force a backend
Authored July 29, 2022
Squashed commit of the following:
Authored August 4, 2022
rollback of https://github.com/google/jax/pull/9596
Authored August 4, 2022
Cache the replacement of `FROM_GDA` to actual shardings present on GDA and check for sharding equality via opsharding.
Authored July 26, 2022
Add `spmd_axis_name` to vmap to allow constraining mapped PartitionSpecs.
Authored August 9, 2022
Refactor BUILD files to avoid individually naming Python dependencies.
Authored August 5, 2022
Try flattening the locals/globals dict in the debugger and have a
Authored August 9, 2022

Browse Other Releases

Latest Pending Unreleased 😎
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
18,442 Ξ”
jax-v0.3.17 Released August 31, 2022
20,878 Ξ”
jax-v0.3.16 Released August 11, 2022
16,659 Ξ”
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,706 Ξ”
jax-v0.3.13 Released May 16, 2022
190 Ξ”
jax-v0.3.12 Released May 16, 2022
13 Ξ”

Top Contributors in jax-v0.3.16

sharadmv
yashk2810
mattjj
hawkinsp
gnecula
jakevdp
LenaMartens
froystig
marcvanzee
NeilGirdhar

Directory Browser for jax-v0.3.16

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.
  • Breaking changes
    • Support for NumPy 1.19 has been dropped, per the deprecation policy. Please upgrade to NumPy 1.20 or newer.
  • Changes
    • Added jax.debug that includes utilities for runtime value debugging such at jax.debug.print and jax.debug.breakpoint.
    • Added new documentation for runtime value debugging
  • Deprecations
    • jax.mask jax.shapecheck APIs have been removed. See #11557.
    • jax.experimental.loops has been removed. See #10278 for an alternative API.
    • jax.tree_util.tree_multimap has been removed. It has been deprecated since JAX release 0.3.5, and jax.tree_util.tree_map is a direct replacement.
    • Removed jax.experimental.stax; it has long been a deprecated alias of jax.example_libraries.stax.
    • Removed jax.experimental.optimizers; it has long been a deprecated alias of jax.example_libraries.optimizers.
    • jax.checkpoint, also known as jax.remat, has a new implementation switched on by default, meaning the old implementation is deprecated; see JEP 11830.
GitClear uses cookies to ensure you get the best experience on our website. Learn more