⬆️ update jax requirement from <=0.9.1,>=0.6.0 to <=0.10.1,>=0.10.1

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.10.1

• New features

  • Added ResizeMethod.AREA to jax.image.resize, which matches TensorFlow's AREA resizing (#20098).
  • Added jax.scipy.linalg.hadamard for constructing Hadamard matrices (#10144).
  • Added jax.scipy.linalg.circulant for constructing circulant matrices (#10144).
  • Added jax.scipy.linalg.dft for constructing discrete Fourier transform matrices (#10144).
  • Added jax.scipy.linalg.leslie for constructing Leslie matrices (#10144).
  • Added jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients (#10144).
  • Added jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices (#10144).
  • Added jax.scipy.linalg.helmert for constructing Helmert matrices (#10144).
  • Moved RNG APIs from "implementations" to dtypes (#27854):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

• Breaking changes

  • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

• Deprecations

  • Passing the copy, order, and ndmin arguments to jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
  • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future version of JAX, this will become an error, if you depend on using them as leaves, pass is_leaf to jax.tree.* methods.

Changelog

Sourced from jax's changelog.

Change log

Best viewed here. For the changes specific to the experimental Pallas APIs, see {ref}pallas-changelog.

JAX follows Effort-based versioning; for a discussion of this and JAX's API compatibility policy, refer to {ref}api-compatibility. For the Python and NumPy version support policy, refer to {ref}version-support-policy.

Unreleased

• New features

  • Added ResizeMethod.AREA to {func}jax.image.resize, which matches TensorFlow's AREA resizing ({jax-issue}[#20098](https://github.com/jax-ml/jax/issues/20098)).
  • Added {func}jax.scipy.linalg.hadamard for constructing Hadamard matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.circulant for constructing circulant matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.dft for constructing discrete Fourier transform matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.leslie for constructing Leslie matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.helmert for constructing Helmert matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Moved RNG APIs from "implementations" to dtypes ({jax-issue}[#27854](https://github.com/jax-ml/jax/issues/27854)):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

• Breaking changes

  • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

• Deprecations

  • Passing the copy, order, and ndmin arguments to {func}jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
  • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future

... (truncated)

Commits

• 619764c Apply patch to revert XLA PR #42396 on release/0.10.1 branch
• 433d98d Prepare for JAX release 0.10.1
• 8c38c32 [Pallas] Enable disabled TPU Interpret Mode test of example kernel
• 6820cf6 [pallas:sc] pl.kernel is now the only API for writing Pallas SC kernels
• 456b22d [state] Fixed a latent bug in _call_primitive_discharge_rule
• d154b80 Update XLA dependency to use revision http://github.com/openxla/xla/commit/9b...
• 11af068 Update XLA dependency to use revision http://github.com/openxla/xla/commit/31...
• 1ecdea1 Improve compilation throughput in multi-threaded auto-tuning environments. Th...
• 46f75a2 Make JaxprEqnContext a bit easier to work with making it easier to add new co...
• 7c01f9b [Pallas MPMD] Refactor tracing context
• Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)

[dependabot]

OK, I won't notify you again about this release, but will get in touch when a new version is available. If you'd rather skip all updates until the next major or minor version, let me know by commenting @dependabot ignore this major version or @dependabot ignore this minor version. You can also ignore all major, minor, or patch releases for a dependency by adding an ignore condition with the desired update_types to your config file.

If you change your mind, just re-open this PR and I'll resolve any conflicts on it.