v0.4.1
Jax release v0.4.1
Added 1
- Introduce jax.Array, a unified array type that subsumes DeviceArray, ShardedDeviceArray, and GlobalDeviceArray types
Changed 6
- jax.Array has been enabled by default and makes some breaking changes to the pjit API
- PartitionSpec and Mesh are now out of experimental with new API endpoints at jax.sharding.PartitionSpec and jax.sharding.Mesh
- with_sharding_constraint's new public endpoint is jax.lax.with_sharding_constraint
- ABSL flag values are no longer read or written after JAX configuration options are initially populated from ABSL flags
- jax2tf.call_tf now uses for TF lowering the first TF device of the same platform as used by the embedding JAX computation
- A number of jax.numpy functions now have their arguments marked as positional-only, matching NumPy
Removed 1
- Support for Python 3.7 has been dropped
Deprecated 2
- jax.experimental.maps.Mesh and jax.experimental.PartitionSpec are deprecated and will be removed in 3 months
- jnp.msort is now deprecated following the deprecation of np.msort in numpy 1.24
- Changes
- Support for Python 3.7 has been dropped, in accordance with JAX's
{ref}
version-support-policy. - We introduce
jax.Arraywhich is a unified array type that subsumesDeviceArray,ShardedDeviceArray, andGlobalDeviceArraytypes in JAX. Thejax.Arraytype helps make parallelism a core feature of JAX, simplifies and unifies JAX internals, and allows us to unifyjitandpjit.jax.Arrayhas been enabled by default in JAX 0.4 and makes some breaking change to thepjitAPI. The jax.Array migration guide can help you migrate your codebase tojax.Array. You can also look at the Distributed arrays and automatic parallelization tutorial to understand the new concepts. PartitionSpecandMeshare now out of experimental. The new API endpoints arejax.sharding.PartitionSpecandjax.sharding.Mesh.jax.experimental.maps.Meshandjax.experimental.PartitionSpecare deprecated and will be removed in 3 months.with_sharding_constraints new public endpoint isjax.lax.with_sharding_constraint.- If using ABSL flags together with
jax.config, the ABSL flag values are no longer read or written after the JAX configuration options are initially populated from the ABSL flags. This change improves performance of readingjax.configoptions, which are used pervasively in JAX. - The jax2tf.call_tf function now uses for TF lowering the first TF device of the same platform as used by the embedding JAX computation. Before, it was using the 0th device for the JAX-default backend.
- A number of
jax.numpyfunctions now have their arguments marked as positional-only, matching NumPy. jnp.msortis now deprecated, following the deprecation ofnp.msortin numpy 1.24. It will be removed in a future release, in accordance with the {ref}api-compatibilitypolicy. It can be replaced withjnp.sort(a, axis=0).
- Support for Python 3.7 has been dropped, in accordance with JAX's
{ref}