Jax

* [GitHub
commits](https://github.com/google/jax/compare/jax-v0.3.5...jax-v0.3.6).
* Changes:
* Upgraded libtpu wheel to a version that fixes a hang when initializing a TPU
pod. Fixes [10218](https://github.com/google/jax/issues/10218).
* Deprecations:
* {mod}`jax.experimental.loops` is being deprecated. See {jax-issue}`10278`
for an alternative API.

* [GitHub
commits](https://github.com/google/jax/compare/jax-v0.3.4...jax-v0.3.5).
* Changes:
* added {func}`jax.random.loggamma` & improved behavior of {func}`jax.random.beta`
and {func}`jax.random.dirichlet` for small parameter values ({jax-issue}`9906`).
* the private `lax_numpy` submodule is no longer exposed in the `jax.numpy` namespace ({jax-issue}`10029`).
* added array creation routines {func}`jax.numpy.frombuffer`, {func}`jax.numpy.fromfunction`,
and {func}`jax.numpy.fromstring` ({jax-issue}`10049`).
* `DeviceArray.copy()` now returns a `DeviceArray` rather than a `np.ndarray` ({jax-issue}`10069`)
* added {func}`jax.scipy.linalg.rsf2csf`
* `jax.experimental.sharded_jit` has been deprecated and will be removed soon.
* Deprecations:
* {func}`jax.nn.normalize` is being deprecated. Use {func}`jax.nn.standardize` instead ({jax-issue}`9899`).
* {func}`jax.tree_util.tree_multimap` is deprecated. Use {func}`jax.tree_util.tree_map` instead ({jax-issue}`5746`).
* `jax.experimental.sharded_jit` is deprecated. Use `pjit` instead.

jaxlib 0.3.5 (April 7, 2022)
* Bug fixes
* Fixed a bug where double-precision complex-to-real IRFFTs would mutate their
input buffers on GPU ({jax-issue}`9946`).
* Fixed incorrect constant-folding of complex scatters ({jax-issue}`10159`)

* [GitHub
commits](https://github.com/google/jax/compare/jax-v0.3.3...jax-v0.3.4).

* [GitHub
commits](https://github.com/google/jax/compare/jax-v0.3.2...jax-v0.3.3).

* [GitHub
commits](https://github.com/google/jax/compare/jax-v0.3.1...jax-v0.3.2).
* Changes:
* The functions `jax.ops.index_update`, `jax.ops.index_add`, which were
deprecated in 0.2.22, have been removed. Please use
[the `.at` property on JAX arrays](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html)
instead, e.g., `x.at[idx].set(y)`.
* Moved `jax.experimental.ann.approx_*_k` into `jax.lax`. These functions are
optimized alternatives to `jax.lax.top_k`.
* {func}`jax.numpy.broadcast_arrays` and {func}`jax.numpy.broadcast_to` now require scalar
or array-like inputs, and will fail if they are passed lists (part of {jax-issue}`7737`).
* The standard jax[tpu] install can now be used with Cloud TPU v4 VMs.
* `pjit` now works on CPU (in addition to previous TPU and GPU support).


jaxlib 0.3.2 (March 16, 2022)
* Changes
* ``XlaComputation.as_hlo_text()`` now supports printing large constants by
passing boolean flag ``print_large_constants=True``.
* Deprecations:
* The ``.block_host_until_ready()`` method on JAX arrays has been deprecated.
Use ``.block_until_ready()`` instead.

* [GitHub
commits](https://github.com/google/jax/compare/jax-v0.3.0...jax-v0.3.1).

* 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 {func}`numpy.testing.assert_allclose()`,
which work directly with JAX arrays ({jax-issue}`9620`).
* `jax.test_util.JaxTestCase` now sets `jax_numpy_rank_promotion='raise'` by default
({jax-issue}`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 {func}`jax.scipy.linalg.schur`, {func}`jax.scipy.linalg.sqrtm`,
{func}`jax.scipy.signal.csd`, {func}`jax.scipy.signal.stft`,
{func}`jax.scipy.signal.welch`.