Jax: jax-v0.7.2 Release

Breaking changes:

• jax.dlpack.from_dlpack no longer accepts a DLPack capsule. This behavior was deprecated and is now removed. The function must be called with an array implementing __dlpack__ and __dlpack_device__.

Changes

• The minimum supported NumPy version is now 2.0. Since SciPy 1.13 is required for NumPy 2.0 support, the minimum supported SciPy version is now 1.13.
• JAX now represents constants in its internal jaxpr representation as a LiteralArray, which is a private JAX type that duck types as a numpy.ndarray. This type may be exposed to users via custom_jvp rules, for example, and may break code that uses isinstance(x, np.ndarray). If this breaks your code, you may convert these arrays to classic NumPy arrays using np.asarray(x).

Bug fixes

• arr.view(dtype=None) now returns the array unchanged, matching NumPy's semantics. Previously it returned the array with a float dtype.
• jax.random.randint now produces a less-biased distribution for 8-bit and 16-bit integer types ({jax-issue}#27742). To restore the previous biased behavior, you may temporarily set the jax_safer_randint configuration to False, but note this is a temporary config that will be removed in a future release.

Deprecations:

• The parameters enable_xla and native_serialization for jax2tf.convert are deprecated and will be removed in a future version of JAX. These were used for jax2tf with non-native serialization, which has been now removed.
• Setting the config state jax_pmap_no_rank_reduction to False is deprecated. By default, jax_pmap_no_rank_reduction will be set to True and jax.pmap shards will not have their rank reduced, keeping the same rank as their enclosing array.