diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index b262282548..5891c36c26 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -12,17 +12,16 @@ jobs: strategy: fail-fast: false matrix: - jax-version: [0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29, - 0.4.30, 0.4.31, 0.4.33, 0.4.34, 0.4.35, 0.4.37, - 0.4.38, 0.5.0, 0.5.1, 0.5.2, 0.5.3, 0.6.0] + jax-version: [0.4.29, 0.4.30, 0.4.31, 0.4.33, 0.4.34, 0.4.35, 0.4.37, + 0.4.38, 0.5.0, 0.5.3, 0.6.0, 0.6.1, 0.6.2, 0.7.2] # 0.4.32 is not available on PyPI - # earlier jax versions are not compatible with other - # dependencies as of 2024-10-04 # 0.4.36 has a bug that causes tests to fail + # 0.5.1 and 0.5.2 installations are broken, see jax#26781 + # 0.7.0 and 0.7.1 have performance issues, see diffrax#680 group: [1, 2] steps: - uses: actions/checkout@v5 - - name: Set up Python 3.10 + - name: Set up Python 3.12 uses: actions/setup-python@v6 with: python-version: '3.12' @@ -35,6 +34,7 @@ jobs: sed -i '1i\jax[cpu] == ${{ matrix.jax-version }}' ./requirements.txt cat ./requirements.txt pip install -r ./devtools/dev-requirements.txt + pip install matplotlib==3.9.2 - name: Verify dependencies run: | python --version @@ -44,7 +44,7 @@ jobs: run: | pwd lscpu - python -m pytest -m unit \ + python -m pytest -v -m unit \ --durations=0 \ --mpl \ --maxfail=1 \ diff --git a/CHANGELOG.md b/CHANGELOG.md index 8cb05e06d8..b6684de73e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,8 +12,10 @@ Bug Fixes - ``desc.compat.rescale`` will now return ``ScaledProfile`` instances for most of its profiles, to fix a bug where improper scaling could occur for certain profile types. Backend +------- - When using any of the ``"proximal-"`` optimization methods, the equilbrium is now always solved before beginning optimization to the specified tolerance (as determined, for example, by ``options={"solve_options":{"ftol"...}}`` passed to the ``desc.optimize.Optimizer.optimize`` call). This ensures the assumptions of the proximal projection method are enforced starting from the first step of the optimization. +- Minimum JAX version bumped up to ``0.4.29`` v0.15.0 diff --git a/desc/batching.py b/desc/batching.py index 3681d58d00..846167948a 100644 --- a/desc/batching.py +++ b/desc/batching.py @@ -2,7 +2,6 @@ from functools import partial -from jax._src import core from jax._src.api import ( _check_input_dtype_jacfwd, _check_input_dtype_jacrev, @@ -21,10 +20,7 @@ _parse_gufunc_signature, _parse_input_dimensions, ) -from jax._src.pjit import auto_axes -from jax._src.sharding_impls import canonicalize_sharding -from jax._src.util import unzip2, wraps -from jax.sharding import PartitionSpec +from jax._src.util import wraps from jax.tree_util import ( tree_flatten, tree_leaves, @@ -36,95 +32,42 @@ from desc.backend import jax, jnp, scan, vmap from desc.utils import errorif -if jax.__version_info__ >= (0, 4, 16): +try: from jax.extend import linear_util as lu -else: +except ImportError: from jax import linear_util as lu -def _scan_leaf(leaf, batch_elems, num_batches, batch_size): - """https://github.com/jax-ml/jax/blob/main/jax/_src/lax/control_flow/loops.py. +try: + from jax._src.lax.control_flow.loops import _batch_and_remainder - References - ---------- - The original copyright notice is as follows - Copyright 2018 The JAX Authors. - Licensed under the Apache License, Version 2.0 (the "License"); - """ - - def f(l): - return l[:batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:]) - - aval = core.typeof(leaf) - if aval.sharding.spec[0] is not None: - raise ValueError( - "0th dimension of leaf passed to `jax.lax.map` should be replicated." - f" Got {aval.str_short(True, True)}" - ) - - out_s = aval.sharding.update( - spec=PartitionSpec(None, None, *aval.sharding.spec[1:]) - ) - out_s = canonicalize_sharding(out_s, "lax.map") - if out_s is not None and out_s.mesh._any_axis_explicit: - return auto_axes(f, out_sharding=out_s, axes=out_s.mesh.explicit_axes)(leaf) - return f(leaf) +except ImportError: + # The old version of JAX doesn't have the required functions and will throw + # an ImportError. We use a simpler version of _batch_and_remainder from an older JAX + # version. + def _batch_and_remainder(x, batch_size: int): + """Taken from JAX 0.5.0. + Function is the same down to JAX 0.4.31. + """ + leaves, treedef = tree_flatten(x) -def _remainder_leaf(leaf, batch_elems): - """https://github.com/jax-ml/jax/blob/main/jax/_src/lax/control_flow/loops.py. - - References - ---------- - The original copyright notice is as follows - Copyright 2018 The JAX Authors. - Licensed under the Apache License, Version 2.0 (the "License"); - """ + scan_leaves = [] + remainder_leaves = [] - def f(l): - return l[batch_elems:] - - sharding = canonicalize_sharding(core.typeof(leaf).sharding, "lax.map") - if sharding is not None and sharding.mesh._any_axis_explicit: - return auto_axes(f, out_sharding=sharding, axes=sharding.mesh.explicit_axes)( - leaf - ) - return f(leaf) - - -def _batch_and_remainder(x, batch_size: int): - """https://github.com/jax-ml/jax/blob/main/jax/_src/lax/control_flow/loops.py. - - References - ---------- - The original copyright notice is as follows - Copyright 2018 The JAX Authors. - Licensed under the Apache License, Version 2.0 (the "License"); - """ - leaves, treedef = tree_flatten(x) - if not leaves: - return x, None - num_batches, remainder = divmod(leaves[0].shape[0], batch_size) - batch_elems = num_batches * batch_size - if num_batches == 0: - remainder_leaves = [_remainder_leaf(leaf, batch_elems) for leaf in leaves] - return None, treedef.unflatten(remainder_leaves) - elif remainder: - scan_leaves, remainder_leaves = unzip2( - [ - ( - _scan_leaf(leaf, batch_elems, num_batches, batch_size), - _remainder_leaf(leaf, batch_elems), + for leaf in leaves: + num_batches = leaf.shape[0] // batch_size + total_batch_elems = num_batches * batch_size + scan_leaves.append( + leaf[:total_batch_elems].reshape( + num_batches, batch_size, *leaf.shape[1:] ) - for leaf in leaves - ] - ) - return treedef.unflatten(scan_leaves), treedef.unflatten(remainder_leaves) - else: - scan_leaves = tuple( - _scan_leaf(leaf, batch_elems, num_batches, batch_size) for leaf in leaves - ) - return treedef.unflatten(scan_leaves), None + ) + remainder_leaves.append(leaf[total_batch_elems:]) + + scan_tree = treedef.unflatten(scan_leaves) + remainder_tree = treedef.unflatten(remainder_leaves) + return scan_tree, remainder_tree def _identity(y): diff --git a/requirements.txt b/requirements.txt index 5fb7f011c0..0d926b32c0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.4.24, != 0.4.36, != 0.5.1, != 0.5.2, <= 0.6.2 +jax >= 0.4.29, != 0.4.36, != 0.5.1, != 0.5.2, <= 0.6.2 colorama <= 0.4.6 diffrax >= 0.4.1, <= 0.7.0 h5py >= 3.0.0, <= 3.14.0 diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 16de8b721c..8b286b6e9f 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1546,13 +1546,14 @@ def test_contravariant_basis_vectors_PEST(eq): @pytest.mark.unit @pytest.mark.slow -@pytest.mark.parametrize("eq", [get("W7-X")]) +@pytest.mark.parametrize("eq", [get("precise_QA")]) def test_PEST_derivative_math(eq): """Verify math to write PEST derivative quantities by redefining θ to θ_PEST.""" from desc.compute import data_index tol = 1e-10 - eq_PEST = eq.to_sfl(4 * eq.L, 5 * eq.M, 4 * eq.N, copy=True, tol=tol) + # TODO: can reduce rtol of test if resolution is increased. See DESC git #1919 + eq_PEST = eq.to_sfl(3 * eq.L, 3 * eq.M, 3 * eq.N, copy=True, tol=tol) keys_DESC = [ "e_theta", @@ -1680,7 +1681,7 @@ def test_PEST_derivative_math(eq): np.testing.assert_allclose( data_to_verify[key_PEST][~near_zero], data[key_DESC][~near_zero], - rtol=3e-3, + rtol=7e-3, err_msg=key_PEST, ) except AssertionError as e: