deps: Update jax

[renovate]

This PR contains the following updates:

Package	Type	Update	Change	Pending
jax	feature-backends	patch	>=0.9.0 → >=0.9.1	0.9.2
jaxlib	feature-cuda-backends	patch	>=0.9.1 → >=0.9.2

---

Release Notes

jax-ml/jax (jax)

v0.9.1

Compare Source

• Changes:

  • JAX tracers that are not of Array type (e.g., of Ref type) will no
    longer report themselves to be instances of Array.
  • Using jax.shard_map in Explicit mode will raise an error
    if the PartitionSpec of input does not match the PartitionSpec specified in
    in_specs. In other words, it will act like an assert instead of an
    implicit reshard.
    in_specs is an optional argument so you can omit specifying it
    and shard_map will infer the PartitionSpec from the argument. If you
    want to reshard your inputs, you can use jax.reshard on the arguments and
    then pass those args to shard_map.

• New features:

  • Added a debug config jax_compilation_cache_check_contents. If set, we miss
    when get() is called on a value that has not been put() by the current
    process, even if the value is actually in the disk cache. When a value is
    put(), we verify that its contents match.

---

Configuration

📅 Schedule: (UTC)

• Branch creation
  • At any time (no schedule defined)
• Automerge
  • At any time (no schedule defined)

🚦 Automerge: Enabled.

♻ Rebasing: Whenever PR becomes conflicted, or you tick the rebase/retry checkbox.

👻 Immortal: This PR will be recreated if closed unmerged. Get config help if that's undesired.

---

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

---

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