deps: Update dependency jax to >=0.9.0

[renovate]

This PR contains the following updates:

Package	Type	Update	Change	Pending
jax	feature-backends	minor	>=0.7.2 → >=0.9.0	0.9.1

---

Release Notes

jax-ml/jax (jax)

v0.9.0

Compare Source

• New features:

  • Added {func}jax.thread_guard, a context manager that detects when devices
    are used by multiple threads in multi-controller JAX.

• Bug fixes:

  • Fixed a workspace size calculation error for pivoted QR (magma_zgeqp3_gpu)
    in MAGMA 2.9.0 when using use_magma=True and pivoting=True.
    ({jax-issue}#34145).

• Deprecations:

  • The flag jax_collectives_common_channel_id was removed.
  • The jax_pmap_no_rank_reduction config state has been removed. The
    no-rank-reduction behavior is now the only supported behavior: a
    jax.pmapped function f sees inputs of the same rank as the input to
    jax.pmap(f). For example, if jax.pmap(f) receives shape (8, 128) on
    8 devices, then f receives shape (1, 128).
  • Setting the jax_pmap_shmap_merge config state is deprecated in JAX v0.9.0
    and will be removed in JAX v0.10.0.
  • {func}jax.numpy.fix is deprecated, anticipating the deprecation of
    {func}numpy.fix in NumPy v2.5.0. {func}jax.numpy.trunc is a drop-in
    replacement.

• Changes:

  • {func}jax.export now supports explicit sharding. This required a new
    export serialization format version that includes the NamedSharding,
    including the abstract mesh, and the partition spec. As part of this
    change we have added a restriction in the use of exported modules: when
    calling them the abstract mesh must match the one used at export time,
    including the axis names. Previously, only the number of the devices
    mattered.

v0.8.2

Compare Source

• Deprecations

  • jax.lax.pvary has been deprecated.
    Please use jax.lax.pcast(..., to='varying') as the replacement.
  • Complex arguments passed to {func}jax.numpy.arange now result in a
    deprecation warning, because the output is poorly-defined.
  • From {mod}jax.core a number of symbols are newly deprecated including:
    call_impl, get_aval, mapped_aval, subjaxprs, set_current_trace,
    take_current_trace, traverse_jaxpr_params, unmapped_aval,
    AbstractToken, and TraceTag.
  • All symbols in {mod}jax.interpreters.pxla are deprecated. These are
    primarily JAX internal APIs, and users should not rely on them.

• Changes:

  • jax's Tracer no longer inherits from jax.Array at runtime. However,
    jax.Array now uses a custom metaclass such isinstance(x, Array) is true
    if an object x represents a traced Array. Only some Tracers represent
    Arrays, so it is not correct for Tracer to inherit from Array.

    For the moment, during Python type checking, we continue to declare Tracer
    as a subclass of Array, however we expect to remove this in a future
    release.

  • jax.experimental.si_vjp has been deleted.
    jax.vjp subsumes it's functionality.

---

Configuration

📅 Schedule: Branch creation - At any time (no schedule defined), Automerge - At any time (no schedule defined).

🚦 Automerge: Enabled.

♻ Rebasing: Whenever PR is behind base branch, or you tick the rebase/retry checkbox.

🔕 Ignore: Close this PR and you won't be reminded about this update again.

---

• If you want to rebase/retry this PR, check this box

---

This PR was generated by Mend Renovate. View the repository job log.