Bump jaxlib from 0.4.28 to 0.4.29

[dependabot]

Bumps jaxlib from 0.4.28 to 0.4.29.

Release notes

Sourced from jaxlib's releases.

Jaxlib release v0.4.29

• Bug fixes

  • Fixed a bug where XLA sharded some concatenation operations incorrectly, which manifested as an incorrect output for cumulative reductions (#21403).
  • Fixed a bug where XLA:CPU miscompiled certain matmul fusions (openxla/xla#13301).
  • Fixes a compiler crash on GPU (google/jax#21396).

• Deprecations

  • jax.tree.map(f, None, non-None) now emits a DeprecationWarning, and will raise an error in a future version of jax. None is only a tree-prefix of itself. To preserve the current behavior, you can ask jax.tree.map to treat None as a leaf value by writing: jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None).

Changelog

Sourced from jaxlib's changelog.

jax 0.4.29 (June 10, 2024)

• Changes

  • We anticipate that this will be the last release of JAX and jaxlib supporting a monolithic CUDA jaxlib. Future releases will use the CUDA plugin jaxlib (e.g. pip install jax[cuda12]).
  • JAX now requires ml_dtypes version 0.4.0 or newer.
  • Removed backwards-compatibility support for old usage of the jax.experimental.export API. It is not possible anymore to use from jax.experimental.export import export, and instead you should use from jax.experimental import export. The removed functionality has been deprecated since 0.4.24.
  • Added is_leaf argument to {func}jax.tree.all & {func}jax.tree_util.tree_all.

• Deprecations

  • jax.sharding.XLACompatibleSharding is deprecated. Please use jax.sharding.Sharding.
  • jax.experimental.Exported.in_shardings has been renamed as jax.experimental.Exported.in_shardings_hlo. Same for out_shardings. The old names will be removed after 3 months.
  • Removed a number of previously-deprecated APIs:
    • from {mod}jax.core: non_negative_dim, DimSize, Shape
    • from {mod}jax.lax: tie_in
    • from {mod}jax.nn: normalize
    • from {mod}jax.interpreters.xla: backend_specific_translations, translations, register_translation, xla_destructure, TranslationRule, TranslationContext, XlaOp.
  • The tol argument of {func}jax.numpy.linalg.matrix_rank is being deprecated and will soon be removed. Use rtol instead.
  • The rcond argument of {func}jax.numpy.linalg.pinv is being deprecated and will soon be removed. Use rtol instead.
  • The deprecated jax.config submodule has been removed. To configure JAX use import jax and then reference the config object via jax.config.
  • {mod}jax.random APIs no longer accept batched keys, where previously some did unintentionally. Going forward, we recommend explicit use of {func}jax.vmap in such cases.

• New Functionality

  • Added {func}jax.experimental.Exported.in_shardings_jax to construct shardings that can be used with the JAX APIs from the HloShardings that are stored in the Exported objects.

jaxlib 0.4.29 (June 10, 2024)

• Bug fixes
  • Fixed a bug where XLA sharded some concatenation operations incorrectly, which manifested as an incorrect output for cumulative reductions (#21403).
  • Fixed a bug where XLA:CPU miscompiled certain matmul fusions (openxla/xla#13301).
  • Fixes a compiler crash on GPU (google/jax#21396).

... (truncated)

Commits

• 991797a Merge pull request #21765 from hawkinsp:release
• 5e7ad60 Removed the double re-exporting of Pallas GPU/TPU APIs
• 3b4039c [Mosaic GPU] Load LLVM lowering interfaces for all dialects
• e071053 Prepare for 0.4.29 release.
• 2ade7e7 [pallas] Move the hardware_generation query in the code path that needs it
• af95803 Merge pull request #21759 from rajasekharporeddy:testbranch1
• 775c6f8 Fix Typos in docs and one error message
• 8fbe65b Update XLA dependency to use revision
• 6617a0d Expand device_put benchmarks to run with different numbers of arrays and in...
• a8246ea Issue a warning where code relies on a bug where treedef.flatten_up_to(...) w...
• 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]

Test Results

0 files   -     3  0 suites   - 3   0s ⏱️ - 54m 52s
0 tests  -   415  0 ✅  -   415  0 💤 ±  0  0 ❌ ±0 
0 runs   - 1 245  0 ✅  - 1 047  0 💤  - 198  0 ❌ ±0 

Results for commit fb587db. ± Comparison against base commit 419ab98.

[github-actions]

🦙 MegaLinter status: ✅ SUCCESS

Descriptor	Linter	Files	Fixed	Errors	Elapsed time
✅ ACTION	actionlint	4		0	0.19s
✅ COPYPASTE	jscpd	yes		no	11.43s
✅ JSON	prettier	2	0	0	0.67s
✅ JSON	v8r	2		0	3.02s
✅ MARKDOWN	markdownlint	1	0	0	0.4s
✅ MARKDOWN	markdown-link-check	1		0	1.31s
✅ MARKDOWN	markdown-table-formatter	1	0	0	0.52s
✅ PYTHON	bandit	231		0	5.66s
✅ PYTHON	black	231	0	0	4.84s
✅ PYTHON	flake8	231		0	3.0s
✅ PYTHON	isort	231	0	0	0.76s
✅ PYTHON	mypy	231		0	6.12s
✅ PYTHON	pylint	231		0	56.1s
✅ PYTHON	ruff	231	0	0	0.07s
✅ REPOSITORY	checkov	yes		no	16.3s
✅ REPOSITORY	gitleaks	yes		no	2.65s
✅ REPOSITORY	git_diff	yes		no	0.04s
✅ REPOSITORY	grype	yes		no	18.15s
✅ REPOSITORY	secretlint	yes		no	6.25s
✅ REPOSITORY	trivy-sbom	yes		no	1.24s
✅ REPOSITORY	trufflehog	yes		no	20.96s
✅ YAML	prettier	8	0	0	0.89s
✅ YAML	v8r	8		0	7.98s
✅ YAML	yamllint	8		0	0.48s

See detailed report in MegaLinter reports

MegaLinter is graciously provided by

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