Jax

jaxlib 0.4.10 (May 11, 2023)

* Changes
* Fixed `'apple-m1' is not a recognized processor for this target (ignoring
processor)` issue that prevented previous release from running on Mac M1.

* Changes
* The flags experimental_cpp_jit, experimental_cpp_pjit and
experimental_cpp_pmap have been removed.
They are now always on.
* Accuracy of singular value decomposition (SVD) on TPU has been improved
(requires jaxlib 0.4.9).

* Deprecations
* `jax.experimental.gda_serialization` is deprecated and has been renamed to
`jax.experimental.array_serialization`.
Please change your imports to use `jax.experimental.array_serialization`.
* The `in_axis_resources` and `out_axis_resources` arguments of pjit have been
deprecated. Please use `in_shardings` and `out_shardings` respectively.
* The function `jax.numpy.msort` has been removed. It has been deprecated since
JAX v0.4.1. Use `jnp.sort(a, axis=0)` instead.
* `in_parts` and `out_parts` arguments have been removed from `jax.xla_computation`
since they were only used with sharded_jit and sharded_jit is long gone.
* `instantiate_const_outputs` argument has been removed from `jax.xla_computation`
since it has been unused for a very long time.

jaxlib 0.4.9 (May 9, 2023)

* Breaking changes
* A major component of the Cloud TPU runtime has been upgraded. This enables
the following new features on Cloud TPU:
* {func}`jax.debug.print`, {func}`jax.debug.callback`, and
{func}`jax.debug.breakpoint()` now work on Cloud TPU
* Automatic TPU memory defragmentation

{func}`jax.experimental.host_callback` is no longer supported on Cloud TPU
with the new runtime component. Please file an issue on the [JAX issue
tracker](https://github.com/google/jax/issues) if the new `jax.debug` APIs
are insufficient for your use case.

The old runtime component will be available for at least the next three
months by setting the environment variable
`JAX_USE_PJRT_C_API_ON_TPU=false`. If you find you need to disable the new
runtime for any reason, please let us know on the [JAX issue
tracker](https://github.com/google/jax/issues).

* Changes
* The minimum jaxlib version has been bumped from 0.4.6 to 0.4.7.

* Deprecations
* CUDA 11.4 support has been dropped. JAX GPU wheels only support
CUDA 11.8 and CUDA 12. Older CUDA versions may work if jaxlib is built
from source.
* `global_arg_shapes` argument of pmap only worked with sharded_jit and has
been removed from pmap. Please migrate to pjit and remove global_arg_shapes
from pmap.

* Changes
* As per https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration
`jax.config.jax_array` cannot be disabled anymore.
* `jax.config.jax_jit_pjit_api_merge` cannot be disabled anymore.
* {func}`jax.experimental.jax2tf.convert` now supports the `native_serialization`
parameter to use JAX's native lowering to StableHLO to obtain a
StableHLO module for the entire JAX function instead of lowering each JAX
primitive to a TensorFlow op. This simplifies the internals and increases
the confidence that what you serialize matches the JAX native semantics.
See [documentation](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md).
As part of this change the config flag `--jax2tf_default_experimental_native_lowering`
has been renamed to `--jax2tf_native_serialization`.
* JAX now depends on `ml_dtypes`, which contains definitions of NumPy types
like bfloat16. These definitions were previously internal to JAX, but have
been split into a separate package to facilitate sharing them with other
projects.
* JAX now requires NumPy 1.21 or newer and SciPy 1.7 or newer.

* Deprecations
* The type `jax.numpy.DeviceArray` is deprecated. Use `jax.Array` instead,
for which it is an alias.
* The type `jax.interpreters.pxla.ShardedDeviceArray` is deprecated. Use
`jax.Array` instead.
* Passing additional arguments to {func}`jax.numpy.ndarray.at` by position is deprecated.
For example, instead of `x.at[i].get(True)`, use `x.at[i].get(indices_are_sorted=True)`
* `jax.interpreters.xla.device_put` is deprecated. Please use `jax.device_put`.
* `jax.interpreters.pxla.device_put` is deprecated. Please use `jax.device_put`.
* `jax.experimental.pjit.FROM_GDA` is deprecated. Please pass in sharded
jax.Arrays as input and remove the `in_shardings` argument to pjit since
it is optional.

jaxlib 0.4.7 (March 27, 2023)

Changes:
* jaxlib now depends on `ml_dtypes`, which contains definitions of NumPy types
like bfloat16. These definitions were previously internal to JAX, but have
been split into a separate package to facilitate sharing them with other
projects.

* Changes
* `jax.tree_util` now contain a set of APIs that allow user to define keys for their
custom pytree node. This includes:
* `tree_flatten_with_path` that flattens a tree and return not only each leaf but
also their key paths.
* `tree_map_with_path` that can map a function that takes the key path as an argument.
* `register_pytree_with_keys` to register how the key path and leaves should looks
like in a custom pytree node.
* `keystr` that pretty-prints a key path.

* {func}`jax2tf.call_tf` has a new parameter `output_shape_dtype` (default `None`)
that can be used to declare the output shape and type of the result. This enables
{func}`jax2tf.call_tf` to work in the presence of shape polymorphism. ({jax-issue}`14734`).

* Deprecations
* The old key-path APIs in `jax.tree_util` are deprecated and will be removed 3 months
from Mar 10 2023:
* `register_keypaths`: use {func}`jax.tree_util.register_pytree_with_keys` instead.
* `AttributeKeyPathEntry` : use `GetAttrKey` instead.
* `GetitemKeyPathEntry` : use `SequenceKey` or `DictKey` instead.

jaxlib 0.4.6 (Mar 9, 2023)

* Deprecations
* `jax.sharding.OpShardingSharding` has been renamed to `jax.sharding.GSPMDSharding`.
`jax.sharding.OpShardingSharding` will be removed in 3 months from Feb 17, 2023.
* The following `jax.Array` methods are deprecated and will be removed 3 months from
Feb 23 2023:
* `jax.Array.broadcast`: use {func}`jax.lax.broadcast` instead.
* `jax.Array.broadcast_in_dim`: use {func}`jax.lax.broadcast_in_dim` instead.
* `jax.Array.split`: use {func}`jax.numpy.split` instead.