Update jax[cpu] requirement from <=0.4.36,>=0.3.2 to >=0.3.2,<=0.4.37

[dependabot]

Updates the requirements on jax[cpu] to permit the latest version.

Release notes

Sourced from jax[cpu]'s releases.

JAX v0.4.37

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 jax.lax.while_loop if the user registers 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.

Changelog

Sourced from jax[cpu]'s changelog.

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 for clues about how to do this). There might also be version skew issues with JAX libraries that do this. If you find this change breaks your non-JAX-internals-using code then try the config.jax_data_dependent_tracing_fallback flag as a workaround, and if you need help updating your code then please file a bug.

  • {func}jax.experimental.jax2tf.convert with native_serialization=False or with enable_xla=False have been deprecated since July 2024, with JAX version 0.4.31. Now we removed support for these use cases. jax2tf with native serialization will still be supported.

  • In jax.interpreters.xla, the xb, xc, and xe symbols have been removed after being deprecated in JAX v0.4.31. Instead use xb = jax.lib.xla_bridge, xc = jax.lib.xla_client, and xe = jax.lib.xla_extension.

  • The deprecated module jax.experimental.export has been removed. It was replaced by {mod}jax.export in JAX v0.4.30. See the migration guide for information on migrating to the new API.

  • The initial argument to {func}jax.nn.softmax and {func}jax.nn.log_softmax has been removed, after being deprecated in v0.4.27.

  • Calling np.asarray on typed PRNG keys (i.e. keys produced by :func:jax.random.key) now raises an error. Previously, this returned a scalar object array.

  • The following deprecated methods and functions in {mod}jax.export have been removed:

    • jax.export.DisabledSafetyCheck.shape_assertions: it had no effect already.
    • jax.export.Exported.lowering_platforms: use platforms.
    • jax.export.Exported.mlir_module_serialization_version: use calling_convention_version.

... (truncated)

Commits

• ffb07cd Update versions for v0.4.37 release.
• 95892fd Use private names for args in api_util to avoid shadowing kwargs keys.
• 65b6088 Avoid index out of range error in carry structure check
• 259194a [Pallas] Fix shard_axis in dma_start interpret mode rule.
• 7e6620a JAX release 0.4.36.
• 23d5c10 [Mosaic:TPU] Fix fully replicated relayout
• 2a4a0e8 [jax:custom_partitioning] Implement SdyShardingRule to support
• f73fa7a Merge pull request #25290 from jakevdp:reduction-where
• a71f9a6 Merge pull request #25271 from jakevdp:fix-vector-norm
• e20a483 [JAX] Add end-to-end execution support in colocated Python API
• 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)

[f0uriest]

@dependabot rebase