Update jax requirement from <0.6.1,>=0.4.25 to >=0.4.25,<0.7.1

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.7.0

• New features:

  • Added jax.P which is an alias for jax.sharding.PartitionSpec.
  • Added jax.tree.reduce_associative.

• Breaking changes:

  • JAX is migrating from GSPMD to Shardy by default. See the migration guide for more information.
  • JAX autodiff is switching to using direct linearization by default (instead of implementing linearization via JVP and partial eval). See migration guide for more information.
  • jax.stages.OutInfo has been replaced with jax.ShapeDtypeStruct.
  • jax.jit now requires fun to be passed by position, and additional arguments to be passed by keyword. Doing otherwise will result in an error starting in v0.7.x. This raised a DeprecationWarning in v0.6.x.
  • The minimum Python version is now 3.11. 3.11 will remain the minimum supported version until July 2026.
  • Layout API renames:
    • Layout, .layout, .input_layouts and .output_layouts have been renamed to Format, .format, .input_formats and .output_formats
    • DeviceLocalLayout, .device_local_layout have been renamed to Layout and .layout
  • jax.experimental.shard module has been deleted and all the APIs have been moved to the jax.sharding endpoint. So use jax.sharding.reshard, jax.sharding.auto_axes and jax.sharding.explicit_axes instead of their experimental endpoints.
  • lax.infeed and lax.outfeed were removed, after being deprecated in JAX 0.6. The transfer_to_infeed and transfer_from_outfeed methods were also removed the Device objects.
  • The jax.extend.core.primitives.pjit_p primitive has been renamed to jit_p, and its name attribute has changed from "pjit" to "jit". This affects the string representations of jaxprs. The same primitive is no longer exported from the jax.experimental.pjit module.
  • The (undocumented) function jax.extend.backend.add_clear_backends_callback has been removed. Users should use jax.extend.backend.register_backend_cache instead.

• Deprecations:

  • {obj}jax.dlpack.SUPPORTED_DTYPES is deprecated; please use the new jax.dlpack.is_supported_dtype function.
  • jax.scipy.special.sph_harm has been deprecated following a similar deprecation in SciPy; use jax.scipy.special.sph_harm_y instead.
  • From {mod}jax.interpreters.xla, the previously deprecated symbols abstractify and pytype_aval_mappings have been removed.
  • jax.interpreters.xla.canonicalize_dtype is deprecated. For canonicalizing dtypes, prefer jax.dtypes.canonicalize_dtype. For checking whether an object is a valid jax input, prefer

... (truncated)

Changelog

Sourced from jax's changelog.

JAX 0.7.0 (July 22, 2025)

• New features:

  • Added jax.P which is an alias for jax.sharding.PartitionSpec.
  • Added {func}jax.tree.reduce_associative.

• Breaking changes:

  • JAX is migrating from GSPMD to Shardy by default. See the migration guide for more information.
  • JAX autodiff is switching to using direct linearization by default (instead of implementing linearization via JVP and partial eval). See migration guide for more information.
  • jax.stages.OutInfo has been replaced with jax.ShapeDtypeStruct.
  • {func}jax.jit now requires fun to be passed by position, and additional arguments to be passed by keyword. Doing otherwise will result in an error starting in v0.7.x. This raised a DeprecationWarning in v0.6.x.
  • The minimum Python version is now 3.11. 3.11 will remain the minimum supported version until July 2026.
  • Layout API renames:
    • Layout, .layout, .input_layouts and .output_layouts have been renamed to Format, .format, .input_formats and .output_formats
    • DeviceLocalLayout, .device_local_layout have been renamed to Layout and .layout
  • jax.experimental.shard module has been deleted and all the APIs have been moved to the jax.sharding endpoint. So use jax.sharding.reshard, jax.sharding.auto_axes and jax.sharding.explicit_axes instead of their experimental endpoints.
  • lax.infeed and lax.outfeed were removed, after being deprecated in JAX 0.6. The transfer_to_infeed and transfer_from_outfeed methods were also removed the Device objects.
  • The jax.extend.core.primitives.pjit_p primitive has been renamed to jit_p, and its name attribute has changed from "pjit" to "jit". This affects the string representations of jaxprs. The same primitive is no longer exported from the jax.experimental.pjit module.
  • The (undocumented) function jax.extend.backend.add_clear_backends_callback has been removed. Users should use jax.extend.backend.register_backend_cache instead.
  • out_sharding arg added to x.at[y].set and x.at[y].add. Previous behavior propagating operand sharding removed. Please use x.at[y].set/add(z, out_sharding=jax.typeof(x).sharding) to retain previous behavior if scatter op requires collectives.

• Deprecations:

  • {obj}jax.dlpack.SUPPORTED_DTYPES is deprecated; please use the new {func}jax.dlpack.is_supported_dtype function.
  • {func}jax.scipy.special.sph_harm has been deprecated following a similar deprecation in SciPy; use {func}jax.scipy.special.sph_harm_y instead.
  • From {mod}jax.interpreters.xla, the previously deprecated symbols

... (truncated)

Commits

• 8bc0ae6 git commit -m "Prepare for JAX release 0.7.0"
• 25a0991 Rollback XLA archival change, current automation does not calculate the shasu...
• 295b92b [Mosaic GPU][NFC] Implement inline_mgpu and custom_primitive for WG seman...
• 603d27e Add visibility to //jax:pretty_printer
• c8e2fb2 Update XLA dependency to use revision
• b6e7105 Fix jax caching bug. The memory space selection logic should take into accoun...
• d16bf37 Merge pull request #30368 from froystig:layout-doc
• 641a768 write docs about controlling array layouts
• 3eaa7f9 Remove SafeNumDevices now that all shardings are implemented in C++
• 2616b4d Loosen test tolerance for `tests/linalg_sharding_test.py::LinalgShardingTest:...
• 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 #114.