# JAX v0.4.34 — JAX v0.4.34

- Product: JAX (https://whatsnew.fyi/product/jax)
- Vendor: Google
- Date: 2024-10-04
- Version: v0.4.34
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.34
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.4.34

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'.

---

- **added** — Wheels for Python 3.13 are now available
- **added** — jax.errors.JaxRuntimeError has been added as a public alias for XlaRuntimeError
- **changed** — jax_pmap_no_rank_reduction flag is now set to True by default
- **changed** — array[0] on a pmap result now introduces a reshape; use array[0:1] instead
- **changed** — Per-shard shape accessed via addressable_shards or addressable_data(0) now has a leading (1, ...) dimension
- **changed** — Default value of --jax_host_callback_legacy configuration is now True, implementing jax.experimental.host_callback APIs in terms of jax.experimental.io_callback
- **deprecated** — Non-arraylike arguments or arraylike arguments with ndim != 1 in jax.numpy.trim_zeros are now deprecated
- **removed** — Internal pretty-printing tools jax.core.pp_* have been removed
- **deprecated** — jax.lib.xla_client.Device is deprecated; use jax.Device instead
- **deprecated** — jax.lib.xla_client.XlaRuntimeError is deprecated; use jax.errors.JaxRuntimeError instead
- **removed** — jax.xla_computation has been deleted
- **removed** — jax.ShapeDtypeStruct no longer accepts the named_shape argument
- **removed** — jax.tree.map(f, None, non-None) now raises an error instead of emitting a DeprecationWarning
- **removed** — jax.sharding.XLACompatibleSharding has been removed; use jax.sharding.Sharding instead
- **fixed** — Fixed a bug where jax.numpy.cumsum would produce incorrect outputs if a non-boolean input was provided and dtype=bool was specified
- **fixed** — Fixed implementation of jax.numpy.ldexp to get correct gradient

* New Functionality
  * This release includes wheels for Python 3.13. Free-threading mode is not yet
    supported.
  * `jax.errors.JaxRuntimeError` has been added as a public alias for the
    formerly private `XlaRuntimeError` type.

* Breaking changes
  * `jax_pmap_no_rank_reduction` flag is set to `True` by default.
    * `array[0]` on a pmap result now introduces a reshape (use `array[0:1]`
      instead).
    * The per-shard shape (accessable via `jax_array.addressable_shards` or
      `jax_array.addressable_data(0))` now has a leading `(1, ...)`. Update code
      that directly accesses shards accordingly. The rank of the per-shard-shape
      now matches that of the global shape which is the same behavior as jit.
      This avoids costly reshapes when passing results from pmap into jit.
  * `jax.experimental.host_callback` has been deprecated since March 2024, with
    JAX version 0.4.26. Now we set the default value of the
    `--jax_host_callback_legacy` configuration value to `True`, which means that
    if your code uses `jax.experimental.host_callback` APIs, those API calls
    will be implemented in terms of the new `jax.experimental.io_callback` API.
    If this breaks your code, for a very limited time, you can set the
    `--jax_host_callback_legacy` to `True`. Soon we will remove that
    configuration option, so you should instead transition to using the
    new JAX callback APIs. See #20385 for a discussion.

* Deprecations
  * In `jax.numpy.trim_zeros`, non-arraylike arguments or arraylike
    arguments with `ndim != 1` are now deprecated, and in the future will result
    in an error.
  * Internal pretty-printing tools `jax.core.pp_*` have been removed, after
    being deprecated in JAX v0.4.30.
  * `jax.lib.xla_client.Device` is deprecated; use `jax.Device` instead.
  * `jax.lib.xla_client.XlaRuntimeError` has been deprecated. Use
    `jax.errors.JaxRuntimeError` instead.

* Deletion:
  * `jax.xla_computation` is deleted. It has been 3 months since its deprecation
    in 0.4.30 JAX release.
    Please use the AOT APIs to get the same functionality as `jax.xla_computation`.
    * `jax.xla_computation(fn)(*args, **kwargs)` can be replaced with
      `jax.jit(fn).lower(*args, **kwargs).compiler_ir('hlo')`.
    * You can also use `.out_info` property of `jax.stages.Lowered` to get the
      output information (like tree structure, shape and dtype).
    * For cross-backend lowering, you can replace
      `jax.xla_computation(fn, backend='tpu')(*args, **kwargs)` with
      `jax.jit(fn).trace(*args, **kwargs).lower(lowering_platforms=('tpu',)).compiler_ir('hlo')`.
  * `jax.ShapeDtypeStruct` no longer accepts the `named_shape` argument.
    The argument was only used by `xmap` which was removed in 0.4.31.
  * `jax.tree.map(f, None, non-None)`, which previously emitted a
    `DeprecationWarning`, now raises an error. `None`
    is only a tree-prefix of itself. To preserve the current behavior, you can
    ask `jax.tree.map` to treat `None` as a leaf value by writing:
    `jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None)`.
  * `jax.sharding.XLACompatibleSharding` has been removed. Please use
    `jax.sharding.Sharding`.

* Bug fixes
  * Fixed a bug where `jax.numpy.cumsum` would produce incorrect outputs
    if a non-boolean input was provided and `dtype=bool` was specified.
  * Edit implementation of `jax.numpy.ldexp` to get correct gradient.