Bump jaxlib from 0.4.30 to 0.4.31

[dependabot]

Bumps jaxlib from 0.4.30 to 0.4.31.

Changelog

Sourced from jaxlib'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).

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)

[github-actions]

🦙 MegaLinter status: ✅ SUCCESS

Descriptor	Linter	Files	Fixed	Errors	Elapsed time
✅ ACTION	actionlint	4		0	0.24s
✅ COPYPASTE	jscpd	yes		no	11.32s
✅ JSON	prettier	2	0	0	0.42s
✅ JSON	v8r	2		0	3.34s
✅ MARKDOWN	markdownlint	1	0	0	0.36s
✅ MARKDOWN	markdown-link-check	1		0	1.28s
✅ MARKDOWN	markdown-table-formatter	1	0	0	0.35s
✅ PYTHON	bandit	231		0	7.63s
✅ PYTHON	black	231	0	0	4.45s
✅ PYTHON	flake8	231		0	4.55s
✅ PYTHON	isort	231	0	0	0.87s
✅ PYTHON	mypy	231		0	6.78s
✅ PYTHON	pylint	231		0	55.42s
✅ PYTHON	ruff	231	0	0	0.1s
✅ REPOSITORY	checkov	yes		no	13.88s
✅ REPOSITORY	gitleaks	yes		no	3.06s
✅ REPOSITORY	git_diff	yes		no	0.02s
✅ REPOSITORY	grype	yes		no	17.38s
✅ REPOSITORY	secretlint	yes		no	4.61s
✅ REPOSITORY	trivy-sbom	yes		no	1.28s
✅ REPOSITORY	trufflehog	yes		no	20.8s
✅ YAML	prettier	8	0	0	0.68s
✅ YAML	v8r	8		0	7.38s
✅ YAML	yamllint	8		0	0.48s

See detailed report in MegaLinter reports

MegaLinter is graciously provided by

[github-actions]

Test Results

    3 files  ±0      3 suites  ±0   57m 49s ⏱️ + 2m 58s
  415 tests ±0    415 ✅ ±0    0 💤 ±0  0 ❌ ±0 
1 245 runs  ±0  1 047 ✅ ±0  198 💤 ±0  0 ❌ ±0 

Results for commit 117d7b5. ± Comparison against base commit 256e0aa.

[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.