JAX release v0.3.18

• GitHub commits.
• Changes
  • Ahead-of-time lowering and compilation functionality (tracked in
    {jax-issue}#7733) is stable and public. See the
    overview and the API docs
    for {mod}jax.stages.
  • Introduced {class}jax.Array, intended to be used for both isinstance checks
    and type annotations for array types in JAX. Notice that this included some subtle
    changes to how isinstance works for {class}jax.numpy.ndarray for jax-internal
    objects, as {class}jax.numpy.ndarray is now a simple alias of {class}jax.Array.
• Breaking changes
  • jax._src is no longer imported into the from the public jax namespace.
    This may break users that were using JAX internals.
  • jax.soft_pmap has been deleted. Please use pjit or xmap instead.
    jax.soft_pmap is undocumented. If it were documented, a deprecation period
    would have been provided.