Update jax requirement from <=0.4.14 to <=0.4.29 in /tests

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.4.29

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

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

Changelog

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

[dependabot]

Superseded by #210.