Update dependency jax to v0.4.13

[renovate]

This PR contains the following updates:

Package	Change	Age	Adoption	Passing	Confidence
jax	0.4.10 -> 0.4.13

---

Release Notes

google/jax

v0.4.13

Compare Source

• Changes

  • jax.jit now allows None to be passed to in_shardings and
    out_shardings. The semantics are as follows:
    • For in_shardings, JAX will mark is as replicated but this behavior
      can change in the future.
    • For out_shardings, we will rely on the XLA GSPMD partitioner to
      determine the output shardings.
  • jax.experimental.pjit.pjit also allows None to be passed to
    in_shardings and out_shardings. The semantics are as follows:
    • If the mesh context manager is not provided, JAX has the freedom to
      choose whatever sharding it wants.
      • For in_shardings, JAX will mark is as replicated but this behavior
        can change in the future.
      • For out_shardings, we will rely on the XLA GSPMD partitioner to
        determine the output shardings.
    • If the mesh context manager is provided, None will imply that the value
      will be replicated on all devices of the mesh.
  • Executable.cost_analysis() works on Cloud TPU
  • Added a warning if a non-allowlisted jaxlib plugin is in use.
  • Added jax.tree_util.tree_leaves_with_path.

• Bug fixes

  • Fixed incorrect wheel name in CUDA 12 releases (#​16362); the correct wheel
    is named cudnn89 instead of cudnn88.

• Deprecations

  • The native_serialization_strict_checks parameter to
    {func}jax.experimental.jax2tf.convert is deprecated in favor of the
    new native_serializaation_disabled_checks ({jax-issue}#16347).

v0.4.12

Compare Source

• Changes

  • Added {class}scipy.spatial.transform.Rotation and {class}scipy.spatial.transform.Slerp

• Deprecations

  • jax.abstract_arrays and its contents are now deprecated. See related
    functionality in :mod:jax.core.
  • jax.numpy.alltrue: use jax.numpy.all. This follows the deprecation
    of numpy.alltrue in NumPy version 1.25.0.
  • jax.numpy.sometrue: use jax.numpy.any. This follows the deprecation
    of numpy.sometrue in NumPy version 1.25.0.
  • jax.numpy.product: use jax.numpy.prod. This follows the deprecation
    of numpy.product in NumPy version 1.25.0.
  • jax.numpy.cumproduct: use jax.numpy.cumprod. This follows the deprecation
    of numpy.cumproduct in NumPy version 1.25.0.
  • jax.sharding.OpShardingSharding has been removed since it has been 3
    months since it was deprecated.

v0.4.11

Compare Source

• Deprecations
  • The following APIs have been removed after a 3 month deprecation period, in
    accordance with the {ref}api-compatibility policy:
    • jax.experimental.PartitionSpec: use jax.sharding.PartitionSpec.
    • jax.experimental.maps.Mesh: use jax.sharding.Mesh
    • jax.experimental.pjit.NamedSharding: use jax.sharding.NamedSharding.
    • jax.experimental.pjit.PartitionSpec: use jax.sharding.PartitionSpec.
    • jax.experimental.pjit.FROM_GDA. Instead pass sharded jax.Array objects
      as input and remove the optional in_shardings argument to pjit.
    • jax.interpreters.pxla.PartitionSpec: use jax.sharding.PartitionSpec.
    • jax.interpreters.pxla.Mesh: use jax.sharding.Mesh
    • jax.interpreters.xla.Buffer: use jax.Array.
    • jax.interpreters.xla.Device: use jax.Device.
    • jax.interpreters.xla.DeviceArray: use jax.Array.
    • jax.interpreters.xla.device_put: use jax.device_put.
    • jax.interpreters.xla.xla_call_p: use jax.experimental.pjit.pjit_p.
    • axis_resources argument of with_sharding_constraint is removed. Please
      use shardings instead.

---

Configuration

📅 Schedule: Branch creation - "before 4am on Monday" (UTC), Automerge - At any time (no schedule defined).

🚦 Automerge: Disabled by config. Please merge this manually once you are satisfied.

♻ Rebasing: Whenever PR becomes conflicted, 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 has been generated by Mend Renovate. View repository job log here.