Bump the jax group across 1 directory with 2 updates

[dependabot]

Updates the requirements on jax and diffrax to permit the latest version.
Updates jax to 0.4.38

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

Updates diffrax to 0.6.2

Release notes

Sourced from diffrax's releases.

Diffrax v0.6.2

• Updated to silence warnings from JAX 0.4.38 (jax.core.Jaxpr -> jax.extend.core.Jaxpr) (Thanks @​jakevdp! #544, #546)
• Added diffrax.ForwardMode for efficient forward-mode autodifferentiation through diffeqsolve. Updated docs to reflect that there are now two main recommended choices, namely ForwardMode for forward-mode and RecursiveCheckpointAdjoint for reverse-mode. (Thanks @​johannahaffner! #537)

Full Changelog: patrick-kidger/diffrax@v0.6.1...v0.6.2

Commits

• 7474749 Bump test versions
• 2589876 Silence warning from pure_callback(vectorized=True)
• 10bff5a Version bump
• da2865e Update documentation on adjoint methods to better emphasise what should and s...
• d28358a Forward mode "adjoint" (#537)
• 3553ae1 Fix pyright failures from numpy 2.
• 613dd0e Import Jaxpr and ClosedJaxpr from jax.exend
• ac9f3ec Fix CI
• d7fc348 Fix docs?
• 78531fa Bump minimum Equinox version to one that is compatible with latest JAX
• 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 <dependency name> major version will close this group update PR and stop Dependabot creating any more for the specific dependency's major version (unless you unignore this specific dependency's major version or upgrade to it yourself)
• @dependabot ignore <dependency name> minor version will close this group update PR and stop Dependabot creating any more for the specific dependency's minor version (unless you unignore this specific dependency's minor version or upgrade to it yourself)
• @dependabot ignore <dependency name> will close this group update PR and stop Dependabot creating any more for the specific dependency (unless you unignore this specific dependency or upgrade to it yourself)
• @dependabot unignore <dependency name> will remove all of the ignore conditions of the specified dependency
• @dependabot unignore <dependency name> <ignore condition> will remove the ignore condition of the specified dependency and ignore conditions

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.64%. Comparing base (6b069f3) to head (c0f722f).
Report is 4 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1501   +/-   ##
=======================================
  Coverage   95.63%   95.64%           
=======================================
  Files         101      101           
  Lines       25542    25542           
=======================================
+ Hits        24428    24429    +1     
+ Misses       1114     1113    -1     

see 1 file with indirect coverage changes

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -4.94 +/- 6.17     | -2.75e-02 +/- 3.44e-02 |  5.30e-01 +/- 2.7e-02  |  5.57e-01 +/- 2.1e-02  |
 test_equilibrium_init_medres            |     -6.57 +/- 6.57     | -2.87e-01 +/- 2.87e-01 |  4.09e+00 +/- 1.1e-01  |  4.37e+00 +/- 2.6e-01  |
 test_equilibrium_init_highres           |     +0.99 +/- 2.92     | +5.40e-02 +/- 1.59e-01 |  5.48e+00 +/- 1.3e-01  |  5.43e+00 +/- 8.6e-02  |
 test_objective_compile_dshape_current   |     -0.00 +/- 1.76     | -1.45e-04 +/- 6.91e-02 |  3.92e+00 +/- 6.4e-02  |  3.92e+00 +/- 2.6e-02  |
 test_objective_compute_dshape_current   |     -2.29 +/- 3.57     | -1.21e-04 +/- 1.88e-04 |  5.15e-03 +/- 8.0e-05  |  5.27e-03 +/- 1.7e-04  |
 test_objective_jac_dshape_current       |     +1.32 +/- 7.26     | +5.61e-04 +/- 3.08e-03 |  4.30e-02 +/- 2.6e-03  |  4.25e-02 +/- 1.7e-03  |
 test_perturb_2                          |     -0.42 +/- 4.90     | -8.53e-02 +/- 9.98e-01 |  2.03e+01 +/- 4.5e-01  |  2.04e+01 +/- 8.9e-01  |
 test_proximal_freeb_jac                 |     +1.00 +/- 2.04     | +7.40e-02 +/- 1.50e-01 |  7.44e+00 +/- 1.3e-01  |  7.36e+00 +/- 6.6e-02  |
 test_solve_fixed_iter                   |     +4.78 +/- 2.01     | +1.53e+00 +/- 6.40e-01 |  3.34e+01 +/- 4.4e-01  |  3.19e+01 +/- 4.7e-01  |
 test_LinearConstraintProjection_build   |     +5.48 +/- 2.96     | +5.61e-01 +/- 3.03e-01 |  1.08e+01 +/- 2.0e-01  |  1.02e+01 +/- 2.2e-01  |
 test_build_transform_fft_midres         |     -0.15 +/- 7.47     | -9.47e-04 +/- 4.71e-02 |  6.29e-01 +/- 3.9e-02  |  6.30e-01 +/- 2.6e-02  |
 test_build_transform_fft_highres        |     +2.01 +/- 4.69     | +2.00e-02 +/- 4.65e-02 |  1.01e+00 +/- 3.7e-02  |  9.93e-01 +/- 2.8e-02  |
 test_equilibrium_init_lowres            |     +1.66 +/- 7.10     | +6.74e-02 +/- 2.88e-01 |  4.13e+00 +/- 2.5e-01  |  4.06e+00 +/- 1.4e-01  |
 test_objective_compile_atf              |     -1.05 +/- 1.36     | -8.75e-02 +/- 1.13e-01 |  8.22e+00 +/- 3.7e-02  |  8.30e+00 +/- 1.1e-01  |
 test_objective_compute_atf              |     +2.77 +/- 5.38     | +4.52e-04 +/- 8.77e-04 |  1.67e-02 +/- 7.8e-04  |  1.63e-02 +/- 4.0e-04  |
 test_objective_jac_atf                  |     +1.30 +/- 2.46     | +2.59e-02 +/- 4.91e-02 |  2.02e+00 +/- 3.9e-02  |  2.00e+00 +/- 2.9e-02  |
 test_perturb_1                          |     -0.78 +/- 2.40     | -1.20e-01 +/- 3.68e-01 |  1.52e+01 +/- 3.4e-01  |  1.54e+01 +/- 1.5e-01  |
 test_proximal_jac_atf                   |     -0.26 +/- 0.91     | -2.18e-02 +/- 7.64e-02 |  8.37e+00 +/- 3.9e-02  |  8.39e+00 +/- 6.5e-02  |
 test_proximal_freeb_compute             |     -1.56 +/- 1.31     | -3.17e-03 +/- 2.67e-03 |  2.01e-01 +/- 2.3e-03  |  2.04e-01 +/- 1.4e-03  |
 test_solve_fixed_iter_compiled          |     +1.31 +/- 4.44     | +2.75e-01 +/- 9.35e-01 |  2.13e+01 +/- 9.1e-01  |  2.11e+01 +/- 1.9e-01  |