Jax: jax-v0.8.1 Release

New features:

• jax.jit now supports the decorator factory pattern; i.e instead of writing python @functools.partial(jax.jit, static_argnames=['n']) def f(x, n): ... you may write python @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:

• jax.sharding.PmapSharding is now deprecated. Please use jax.NamedSharding instead.
• jx.device_put_replicated is now deprecated. Please use jax.device_put with the appropriate sharding instead.
• jax.device_put_sharded is now deprecated. Please use jax.device_put with the appropriate sharding instead.
• Default axis_types of jax.make_mesh will change in JAX v0.9.0 to return jax.sharding.AxisType.Explicit. Leaving axis_types unspecified will raise a DeprecationWarning.
• jax.cloud_tpu_init and its contents were deprecated. There is no reason for a user to import or use the contents of this module; JAX handles this for you automatically if needed.