Adds missing .pyi files that were missing from the previous release ({jax-issue}[#12536](https://github.com/google/jax/issues/12536)).
Fixes an incompatibility between jax 0.3.19 and the libtpu version it pinned ({jax-issue}[#12550](https://github.com/google/jax/issues/12550)). Requires jaxlib 0.3.20.
Fix incorrect pip url in setup.py comment ({jax-issue}[#12528](https://github.com/google/jax/issues/12528)).
jaxlib 0.3.20 (Sep 28, 2022)
Bug fixes
Fixes support for limiting the visible CUDA devices via
jax_cuda_visible_devices in distributed jobs. This functionality is needed for
the JAX/SLURM integration on GPU ({jax-issue}[#12533](https://github.com/google/jax/issues/12533)).
Ahead-of-time lowering and compilation functionality (tracked in
{jax-issue}[#7733](https://github.com/google/jax/issues/7733)) is stable and public. See the
overview and the API docs
for {mod}jax.stages.
Introduced {class}jax.Array, intended to be used for both isinstance checks
and type annotations for array types in JAX. Notice that this included some subtle
changes to how isinstance works for {class}jax.numpy.ndarray for jax-internal
objects, as {class}jax.numpy.ndarray is now a simple alias of {class}jax.Array.
Breaking changes
jax._src is no longer imported into the from the public jax namespace.
This may break users that were using JAX internals.
jax.soft_pmap has been deleted. Please use pjit or xmap instead.
jax.soft_pmap is undocumented. If it were documented, a deprecation period
would have been provided.
Fix corner case issue in gradient of lax.pow with an exponent of zero
({jax-issue}12041)
Breaking changes
{func}jax.checkpoint, also known as {func}jax.remat, no longer supports
the concrete option, following the previous version's deprecation; see
JEP 11830.
Changes
Added {func}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 {func}jax.numpy.tile
({jax-issue}[#11944](https://github.com/google/jax/issues/11944)).
DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead.
jax 0.3.16
... (truncated)
Commits
0282b4b Merge pull request #12538 from jakevdp:bundle-pyi
aafc70d Merge pull request #12556 from hawkinsp:rocm
5fe7a54 Merge pull request #12555 from hawkinsp:release
39eabe8 Merge pull request #12552 from hawkinsp:nccl
f7bafb3 Disable multiprocess_gpu_test that fails on ROCm.
Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot merge` will merge this PR after your CI passes on it
- `@dependabot squash and merge` will squash and merge this PR after your CI passes on it
- `@dependabot cancel merge` will cancel a previously requested merge and block automerging
- `@dependabot reopen` will reopen this PR if it is closed
- `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Updates the requirements on jax to permit the latest version.
Release notes
Sourced from jax's releases.
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
0282b4b
Merge pull request #12538 from jakevdp:bundle-pyiaafc70d
Merge pull request #12556 from hawkinsp:rocm5fe7a54
Merge pull request #12555 from hawkinsp:release39eabe8
Merge pull request #12552 from hawkinsp:ncclf7bafb3
Disable multiprocess_gpu_test that fails on ROCm.8d86436
jax/jaxlib 0.3.20 release candidate.eabb91e
Fix test failure in GPU CI if NCCL_DEBUG is enabled.96abd9a
Merge pull request #12540 from sharadmv:cond-lowering-fix96a85bd
Make addressable_shards a property like local_shards9489068
Merge pull request #12546 from mattjj:issue12542Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting
@dependabot rebase
.Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot merge` will merge this PR after your CI passes on it - `@dependabot squash and merge` will squash and merge this PR after your CI passes on it - `@dependabot cancel merge` will cancel a previously requested merge and block automerging - `@dependabot reopen` will reopen this PR if it is closed - `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)