# JAX v0.4.1 — Jax release v0.4.1

- Product: JAX (https://whatsnew.fyi/product/jax)
- Vendor: Google
- Date: 2022-12-13
- Version: v0.4.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.4.1

What's New is an index, not a publisher: every entry below links to the vendor's own release notes, which are the authoritative source. Entries are labelled where they are hand-curated sample data, pre-releases, or drawn from a secondary source such as a developer blog.

Reuse: the summaries, labels and curation here are © What's New. Quote freely with attribution and a link back; wholesale republication of the corpus is not permitted — terms: https://whatsnew.fyi/terms. The vendors' own release notes remain their publishers'.

---

- **removed** — Support for Python 3.7 has been dropped
- **added** — Introduce jax.Array, a unified array type that subsumes DeviceArray, ShardedDeviceArray, and GlobalDeviceArray types
- **changed** — jax.Array has been enabled by default and makes some breaking changes to the pjit API
- **changed** — PartitionSpec and Mesh are now out of experimental with new API endpoints at jax.sharding.PartitionSpec and jax.sharding.Mesh
- **deprecated** — jax.experimental.maps.Mesh and jax.experimental.PartitionSpec are deprecated and will be removed in 3 months
- **changed** — with_sharding_constraint's new public endpoint is jax.lax.with_sharding_constraint
- **changed** — ABSL flag values are no longer read or written after JAX configuration options are initially populated from ABSL flags
- **changed** — jax2tf.call_tf now uses for TF lowering the first TF device of the same platform as used by the embedding JAX computation
- **changed** — A number of jax.numpy functions now have their arguments marked as positional-only, matching NumPy
- **deprecated** — 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.Array` which is a unified array type that subsumes
    `DeviceArray`, `ShardedDeviceArray`, and `GlobalDeviceArray` types in JAX.
    The `jax.Array` type helps make parallelism a core feature of JAX,
    simplifies and unifies JAX internals, and allows us to unify `jit` and
    `pjit`.  `jax.Array` has been enabled by default in JAX 0.4 and makes some
    breaking change to the `pjit` API.  The [jax.Array migration
    guide](https://jax.readthedocs.io/en/latest/jax_array_migration.html) can
    help you migrate your codebase to `jax.Array`. You can also look at the
    [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)
    tutorial to understand the new concepts.
  * `PartitionSpec` and `Mesh` are now out of experimental. The new API endpoints
    are `jax.sharding.PartitionSpec` and `jax.sharding.Mesh`.
    `jax.experimental.maps.Mesh` and `jax.experimental.PartitionSpec` are
    deprecated and will be removed in 3 months.
  * `with_sharding_constraint`s new public endpoint is
    `jax.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 reading
    `jax.config` options, 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.numpy` functions now have their arguments marked as
    positional-only, matching NumPy.
  * `jnp.msort` is now deprecated, following the deprecation of `np.msort` in numpy 1.24.
    It will be removed in a future release, in accordance with the {ref}`api-compatibility`
    policy. It can be replaced with `jnp.sort(a, axis=0)`.