●Change8

jax-v0.10.1

Breaking Changes
📦 jaxView on GitHub →
⚠ 1 breaking✨ 11 features⚡ 2 deprecations🔧 13 symbols

Summary

This release introduces several new linear algebra functions in `jax.scipy.linalg` and updates RNG API handling by moving related functionality to dtypes. It also deprecates positional arguments for certain array creation parameters and modifies the mesh context manager usage.

⚠️ Breaking Changes

  • The `with mesh:` context manager is deprecated. Users must switch to using `with jax.set_mesh(mesh):` instead.

Migration Steps

  1. Replace usage of the deprecated `with mesh:` context manager with `with jax.set_mesh(mesh):`.

✨ New Features

  • Added `jax.image.resize` support for `ResizeMethod.AREA`, aligning with TensorFlow's behavior.
  • Added `jax.scipy.linalg.hadamard` for constructing Hadamard matrices.
  • Added `jax.scipy.linalg.circulant` for constructing circulant matrices.
  • Added `jax.scipy.linalg.dft` for constructing discrete Fourier transform matrices.
  • Added `jax.scipy.linalg.leslie` for constructing Leslie matrices.
  • Added `jax.scipy.linalg.companion` for constructing companion matrices from polynomial coefficients.
  • Added `jax.scipy.linalg.fiedler` for constructing symmetric Fiedler matrices.
  • Added `jax.scipy.linalg.helmert` for constructing Helmert matrices.
  • RNG APIs were moved from "implementations" to dtypes.
  • Added `jax.random.key_dtype` to retrieve the dtype corresponding to a PRNG implementation name.
  • `jax.random.key` and `wrap_key_data` now accept an optional `dtype` argument.

Affected Symbols

jax.image.resizejax.scipy.linalg.hadamardjax.scipy.linalg.circulantjax.scipy.linalg.dftjax.scipy.linalg.lesliejax.scipy.linalg.companionjax.scipy.linalg.fiedlerjax.scipy.linalg.helmertjax.random.key_dtypejax.random.keywrap_key_datajax.numpy.arrayjax.tree.*

⚡ Deprecations

  • Passing `copy`, `order`, and `ndmin` arguments positionally to `jax.numpy.array` is deprecated; they must now be passed as keyword arguments.
  • Python `dict_values`, generators, zip return type, and general iterators are deprecated by default when used as leaves in pytrees. Users relying on this behavior should pass an `is_leaf` argument to `jax.tree.*` methods.