New features:

• jax.jit now supports the decorator factory pattern; i.e instead of writing

  @functools.partial(jax.jit, static_argnames=['n'])
  def f(x, n):
    ...
  

  you may write

  @jax.jit(static_argnames=['n'])
  def f(x, n):
    ...
  

Changes:

• jax.lax.linalg.eigh now accepts an implementation argument to select between QR (CPU/GPU), Jacobi (GPU/TPU), and QDWH (TPU) implementations. The EighImplementation enum is publicly exported from jax.lax.linalg.

• jax.lax.linalg.svd now implements an algorithm that uses the polar decomposition on CUDA GPUs. This is also an alias for the existing algorithm on TPUs.

Bug fixes:

• Fixed a bug introduced in JAX 0.7.2 where eigh failed for large matrices on GPU (#33062).

Deprecations: