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

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.4.38

• Changes:

  • jax.tree.flatten_with_path and jax.tree.map_with_path are added as shortcuts of the corresponding tree_util functions.

• Deprecations

  • a number of APIs in the internal jax.core namespace have been deprecated. Most were no-ops, were little-used, or can be replaced by APIs of the same name in jax.extend.core; see the documentation for {mod}jax.extend for information on the compatibility guarantees of these semi-public extensions.
  • Several previously-deprecated APIs have been removed, including:
    • from jax.core: check_eqn, check_type, check_valid_jaxtype, and non_negative_dim.
    • from jax.lib.xla_bridge: xla_client and default_backend.
    • from jax.lib.xla_client: _xla and bfloat16.
    • from jax.numpy: round_.

• New Features

  • jax.export.export can be used for device-polymorphic export with shardings constructed with {func}jax.sharding.AbstractMesh. See the jax.export documentation.
  • Added jax.lax.split. This is a primitive version of jax.numpy.split, added because it yields a more compact transpose during automatic differentiation.

Changelog

Sourced from jax's changelog.

jax 0.4.38 (Dec 17, 2024)

• Changes:

  • jax.tree.flatten_with_path and jax.tree.map_with_path are added as shortcuts of the corresponding tree_util functions.

• Deprecations

  • a number of APIs in the internal jax.core namespace have been deprecated. Most were no-ops, were little-used, or can be replaced by APIs of the same name in {mod}jax.extend.core; see the documentation for {mod}jax.extend for information on the compatibility guarantees of these semi-public extensions.
  • Several previously-deprecated APIs have been removed, including:
    • from {mod}jax.core: check_eqn, check_type, check_valid_jaxtype, and non_negative_dim.
    • from {mod}jax.lib.xla_bridge: xla_client and default_backend.
    • from {mod}jax.lib.xla_client: _xla and bfloat16.
    • from {mod}jax.numpy: round_.

• New Features

  • {func}jax.export.export can be used for device-polymorphic export with shardings constructed with {func}jax.sharding.AbstractMesh. See the jax.export documentation.
  • Added {func}jax.lax.split. This is a primitive version of {func}jax.numpy.split, added because it yields a more compact transpose during automatic differentiation.

jax 0.4.37 (Dec 9, 2024)

This is a patch release of jax 0.4.36. Only "jax" was released at this version.

• Bug fixes
  • Fixed a bug where jit would error if an argument was named f (#25329).
  • Fix a bug that will throw index out of range error in {func}jax.lax.while_loop if the user register pytree node class with different aux data for the flatten and flatten_with_path.
  • Pinned a new libtpu release (0.0.6) that fixes a compiler bug on TPU v6e.

jax 0.4.36 (Dec 5, 2024)

• Breaking Changes

  • This release lands "stackless", an internal change to JAX's tracing machinery. We made trace dispatch purely a function of context rather than a function of both context and data. This let us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind, and so on. The change should only affect users that use JAX internals.

    If you do use JAX internals then you may need to update your code (see jax-ml/jax@c36e1f7

... (truncated)

Commits

• c4e3f5a Release v0.4.38
• 0fa5419 Merge pull request #25456 from jakevdp:xla-abstractify
• 7fe2579 Update XLA dependency to use revision
• 7bbba40 Merge pull request #25473 from adamjstewart:build/no-git
• 7de9eb2 Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
• 4911a39 [Mosaic TPU] Add support for the interleaved pack format to tpu.unpack_subele...
• 36b12d5 [Mosaic GPU] Add end-to-end lowering example for a pointwise kernel using the...
• 2c722d9 Cleanup: toward merging core.concrete_aval & xla.abstractify
• 473e2bf Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr...
• e1f037b Reduce the max_examples for splash attention tests
• 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)