v0.3.18
JAX release v0.3.18
Added 2
- Ahead-of-time lowering and compilation functionality is now stable and public via jax.stages
- Introduce jax.Array for isinstance checks and type annotations for array types in JAX
Changed 1
- jax.numpy.ndarray is now a simple alias of jax.Array, changing isinstance behavior for jax-internal objects
Removed 2
- jax._src is no longer imported into the public jax namespace
- jax.soft_pmap has been deleted; use pjit or xmap instead
- 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 bothisinstancechecks and type annotations for array types in JAX. Notice that this included some subtle changes to howisinstanceworks for {class}jax.numpy.ndarrayfor jax-internal objects, as {class}jax.numpy.ndarrayis now a simple alias of {class}jax.Array.
- Ahead-of-time lowering and compilation functionality (tracked in
{jax-issue}
- Breaking changes
jax._srcis no longer imported into the from the publicjaxnamespace. This may break users that were using JAX internals.jax.soft_pmaphas been deleted. Please usepjitorxmapinstead.jax.soft_pmapis undocumented. If it were documented, a deprecation period would have been provided.