jax-ml/jax jax-v0.2.25

JAX release v0.2.25

on GitHub

New features:

• (Experimental) jax.distributed.initialize exposes multi-host GPU backend.
• jax.random.permutation supports new independent keyword argument
  ({jax-issue}#8430)

Breaking changes

• Moved jax.experimental.stax to jax.example_libraries.stax
• Moved jax.experimental.optimizers to jax.example_libraries.optimizers

New features:

• Added jax.lax.linalg.qdwh.