Sparse pullback for big performance gain

CHANGELOG.md

@@ -7,7 +7,6 @@ New Features
 - Sub-objectives of an `ObjectiveFunction` can now have different `use_jit` values than the `ObjectiveFunction`. These objectives have to be built before building the `ObjectiveFunction`.
 - Adds ``num_neighbors`` parameter to ``CoilSetMinDistance`` that limits the pairwise distance computation to the nearest neighbors per coil, reducing memory useage for large coilsets.
 - Method to plot frequency spectrum of inverse stream map in field line coordinates ``Bounce2D.plot_angle_spectrum``.
-- Method to compute bounce integrals in batches is now added to the public API ``Bounce2D.batch``.
 - Initiated deprecation of ``Bounce2D.compute_fieldline_length`` in favor of ``eq.compute("V_psi")``.
 - The quadrature resolution in ``Bounce2D.compute_fieldline_length`` now corresponds to the resolution over a single field period instead of the resolution over a toroidal transit.
 - Adds an optional attribute `ion_density` to the `Equilibrium` class, to allow the ion density profile to be set independently of the electron density and effective atomic number.
@@ -17,20 +16,29 @@ New Features
 Bug Fixes
 
 - Fixes SyntaxError thrown when loading hdf5 data from file-like objects.
+- Fixes ``pitch_batch_size`` argument getting ignored in compute functions.
 - Fixes a bug in `OmnigenousField.change_resolution` when changing `L_B`.
 - Scaling a `ScaledProfile` or taking power of a `PowerProfile` now only updates the `scale`/`power` attributes instead of nesting the `ScaledProfile`/`PowerProfile`s.
 - `jax.Array`s in `_static_attrs` will be automatically converted to `np.ndarray` to prevent stalling code. In general, jax arrays should be omitted in `_static_attrs`.
 
 Performance Improvements
 
+- Sparse reverse-mode differentiation was introduced to DESC
+  to yield significant performance improvements [#2170](https://github.com/PlasmaControl/DESC/pull/2170).
+  Plumbing to use this method was added to DESC that
+  will be progressively taken advantage of in the future.
 - Reduces import time of `desc` modules.
     - Now, `desc.compute._build_data_index` uses depth-first search algorithm to construct the dependency tree.
     - Some of the default value computations at import time are removed (i.e. `desc.integrals.bounce_integral.default_quad`)
 - [Significantly improves convergence of inverse stream maps](https://github.com/PlasmaControl/DESC/pull/1919).
-- Check-pointing to bounce integrals to improve speed and reduce memory of reverse mode differentiation.
 - Resolves a JAX memory regression in bounce integrals by avoiding materialization of a large tensor in memory. Previously, we had closed the issue by adding nuffts as a workaround. This update actually solves the issue for the case when a user specifies to not use nuffts as well.
 - ``ObjectiveFunction.print_value`` can now use the previously computed ``compute_scaled_error`` values to print. For bounded objectives, we fall back to computing ``compute_unscaled``. Additionally, ``compute_scaled_error`` and array splitting are used in other parts of the code to prevent recompilation for one-time tasks, which makes initialization faster.
 
+Breaking Changes
+
+- The parameter ``num_transit`` in ``EffectiveRipple``, ``Gamma_c``, ``Bounce2D`` and related functions has been changed to ``num_field_periods``. This should make using a consistent resolution across different equilibria easier.
+- The parameter ``Y_B`` in ``EffectiveRipple``, ``Gamma_c``, ``Bounce2D`` is now the resolution over a single field period rather than a full toroidal transit. This should make using a consistent resolution across different equilibria easier.
+
 Deprecations
 
 - `constants` argument of `compute`, `jvp`, `jac`, `grad` and `hess` methods (including all of their variants) to all objective classes (including `ObjectiveFunction` and wrappers) is deprecated and will be removed in a future release. This argument was not necessary, and the code will still work if user doesn't pass it. Users should update their custom objectives for this change. In addition, `constants` property of the `ObjectiveFunction` and all sub-classes of `_Objective` is deprecated.

desc/batching.py

@@ -197,6 +197,21 @@ def vmap_chunked(
     - https://github.com/jax-ml/jax/issues/26689
     - https://github.com/jax-ml/jax/issues/27591
     - https://github.com/jax-ml/jax/issues/31919
+    - Due to an actively worked on issue in JAX,
+      https://docs.jax.dev/en/latest/jep/
+      2026-custom-derivatives.html#main-problem-descriptions,
+      this function can simply ignore custom derivative rules
+      of the function in wraps if ``chunk_size`` is not ``None``,
+      and therefore can damp the effeciency gains of ``sparse_pullback``.
+      Use ``batch_map`` instead to avoid this,
+      or try to make a hack with jax.custom_transforms to bypass this.
+    - Only out axes = 0 is supported.
+
+    See Also
+    --------
+    batch_map
+        If the function supports native vectorization, use ``batch_map`` instead
+        for the reasons discussed the docstring.
 
     Parameters
     ----------
@@ -231,20 +246,45 @@ def vmap_chunked(
 
 
 def batch_map(
-    fun, fun_input, /, batch_size=None, *, reduction=None, chunk_reduction=identity
+    fun,
+    fun_input,
+    /,
+    batch_size=None,
+    *,
+    reduction=None,
+    chunk_reduction=identity,
+    strip_dim0=False,
 ):
     """Compute ``chunk_reduction(fun(fun_input))`` in batches.
 
-    This utility is like ``vmap_chunked`` except that ``fun`` is assumed to be
-    vectorized natively. No JAX vectorization such as ``vmap`` is applied to the
-    supplied function. This makes compilation faster and avoids the weaknesses of
-    applying JAX vectorization, such as executing all branches of code conditioned on
-    dynamic values. For example, this function would be useful for GitHub issue #1303.
+    Notes
+    -----
+    This method does not automatically wrap ``fun`` with ``vmap``.
+    Unless ``fun`` is already wrapped with ``vmap``, the leading dimension
+    of ``fun_input`` will not be stripped before it is passed into ``fun``.
+    This can be inconvenient for nesting calls to ``batch_map``,
+    since only batching along the first axis is supported.
+    However, the ``strip_dim0`` flag should cover the most common case
+    of nesting calls where ``batch_size`` is one on the outermost call.
+
+    If ``fun`` is natively vectorized, this can be preferable to ``vmap_chunked``
+    to reduce compilation time, avoid issues such as executing all branches of
+    code conditioned on dynamic values, or avoid messing up the behavior of
+    jvp's and vjp's under vmap, e.g.
+    https://docs.jax.dev/en/latest/jep/
+    2026-custom-derivatives.html#main-problem-descriptions.
+
+    Only out axes = 0 is supported.
+
+    See Also
+    --------
+    vmap_chunked
+        If the function does not support native vectorization.
 
     Parameters
     ----------
     fun : callable
-        Natively vectorized function.
+        Vectorized function.
     fun_input : pytree
         Data to split into batches to feed to ``fun``.
     batch_size : int or None
@@ -257,19 +297,30 @@ def batch_map(
         Chunk-wise reduction operation.
         Should typically apply ``reduction`` along the mapped axis,
         e.g. ``jnp.add.reduce``.
+    strip_dim0 : bool
+        Whether to strip the leading dim of ``fun_input`` before passing it
+        to ``fun``; see notes. This flag only works if ``batch_size`` is one.
+        It should be set to ``False`` if ``fun`` is wrapped in ``vmap``.
+        Default is ``False``.
 
     Returns
     -------
     fun_output
         Returns ``chunk_reduction(fun(fun_input))``.
 
     """
-    return (
-        chunk_reduction(fun(fun_input))
-        if batch_size is None
-        else _evaluate_in_chunks(
-            fun, batch_size, (0,), reduction, chunk_reduction, fun_input
-        )
+    if batch_size is None:
+        return chunk_reduction(fun(fun_input))
+    if strip_dim0 and batch_size == 1:
+        return _scanmap(fun, 0, reduction, identity)(fun_input)
+
+    return _evaluate_in_chunks(
+        fun,
+        batch_size,
+        (0,),
+        reduction,
+        chunk_reduction,
+        fun_input,
     )