Jax

* Changes
* Add `JAX_PLATFORMS=tpu,cpu` as default setting in TPU initialization,
so JAX will raise an error if TPU cannot be initialized instead of falling
back to CPU. Set `JAX_PLATFORMS=''` to override this behavior and automatically
choose an available backend (the original default), or set `JAX_PLATFORMS=cpu`
to always use CPU regardless of if the TPU is available.
* Deprecations
* Several test utilities deprecated in JAX v0.3.8 are now removed from
{mod}`jax.test_util`.

jaxlib 0.3.22 (Oct 11, 2022)

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.20...jax-v0.3.21).
* Changes
* The persistent compilation cache will now warn instead of raising an
exception on error ({jax-issue}`12582`), so program execution can continue
if something goes wrong with the cache. Set
`JAX_RAISE_PERSISTENT_CACHE_ERRORS=true` to revert this behavior.

* Bug fixes:
* Adds missing `.pyi` files that were missing from the previous release ({jax-issue}`12536`).
* Fixes an incompatibility between `jax` 0.3.19 and the libtpu version it pinned ({jax-issue}`12550`). Requires jaxlib 0.3.20.
* Fix incorrect `pip` url in `setup.py` comment ({jax-issue}`12528`).

jaxlib 0.3.20 (Sep 28, 2022)
* [GitHub commits](https://github.com/google/jax/compare/jaxlib-v0.3.15...jaxlib-v0.3.20).
* 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`).

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.18...jax-v0.3.19).
* Fixes required jaxlib version.

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.17...jax-v0.3.18).
* Changes
* Ahead-of-time lowering and compilation functionality (tracked in
{jax-issue}`7733`) is stable and public. See [the
overview](https://jax.readthedocs.io/en/latest/aot.html) 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 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.

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.16...jax-v0.3.17).
* Bugs
* 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](https://jax.readthedocs.io/en/latest/jep/11830-new-remat-checkpoint.html).
* 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`).
* `DeviceArray.to_py()` has been deprecated. Use `np.asarray(x)` instead.