Update jax requirement from <=0.4.24 to <=0.4.31

[dependabot]

Updates the requirements on jax to permit the latest version.

Changelog

Sourced from jax's changelog.

jax 0.4.31 (July 29, 2024)

• Deletion

  • xmap has been deleted. Please use {func}shard_map as the replacement.

• Changes

  • The minimum CuDNN version is v9.1. This was true in previous releases also, but we now declare this version constraint formally.
  • The minimum Python version is now 3.10. 3.10 will remain the minimum supported version until July 2025.
  • The minimum NumPy version is now 1.24. NumPy 1.24 will remain the minimum supported version until December 2024.
  • The minimum SciPy version is now 1.10. SciPy 1.10 will remain the minimum supported version until January 2025.
  • {func}jax.numpy.ceil, {func}jax.numpy.floor and {func}jax.numpy.trunc now return the output of the same dtype as the input, i.e. no longer upcast integer or boolean inputs to floating point.
  • libdevice.10.bc is no longer bundled with CUDA wheels. It must be installed either as a part of local CUDA installation, or via NVIDIA's CUDA pip wheels.
  • {class}jax.experimental.pallas.BlockSpec now expects block_shape to be passed before index_map. The old argument order is deprecated and will be removed in a future release.
  • Updated the repr of gpu devices to be more consistent with TPUs/CPUs. For example, cuda(id=0) will now be CudaDevice(id=0).
  • Added the device property and to_device method to {class}jax.Array, as part of JAX's Array API support.

• Deprecations

  • Removed a number of previously-deprecated internal APIs related to polymorphic shapes. From {mod}jax.core: removed canonicalize_shape, dimension_as_value, definitely_equal, and symbolic_equal_dim.
  • HLO lowering rules should no longer wrap singleton ir.Values in tuples. Instead, return singleton ir.Values unwrapped. Support for wrapped values will be removed in a future version of JAX.
  • {func}jax.experimental.jax2tf.convert with native_serialization=False or enable_xla=False is now deprecated and this support will be removed in a future version. Native serialization has been the default since JAX 0.4.16 (September 2023).
  • The previously-deprecated function jax.random.shuffle has been removed; instead use jax.random.permutation with independent=True.

jaxlib 0.4.31 (July 29, 2024)

• Bug fixes
  • Fixed a bug that meant that negative static_argnums to a jit were mishandled by the jit dispatch fast path.
  • Fixed a bug that meant triangular solves of batches of singular matrices produce nonsensical finite values, instead of inf or nan (#3589, #15429).

jax 0.4.30 (June 18, 2024)

... (truncated)

Commits

• 7fd9302 Start JAX and jaxlib 0.4.31 release
• f070c06 Merge pull request #22703 from Rifur13:plugin-fix
• 9beb4f1 Merge pull request #19760 from Blair-Johnson:fix-pytree-grads-sparse
• 802a14c Re-pack gradients of jax.experimental.sparse.grad() to match original pytrees...
• 85e83b5 Merge pull request #22690 from jakevdp:inplace-doc
• d1c0d99 Bump the minimum CUDNN version to v9.1.
• 6127baa Ignore the Deprecation warning produced about native_serialization=False.
• fd23b87 Bump minimum SciPy version to 1.10.
• cfa1e78 Improve documentation for jnp.put, jnp.place, jnp.fill_diagonal
• 0224235 Skip cuda backend initialization if no nvidia GPUs are visible.
• 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 merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @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)