Add warning for batched mode when a sub objective has rev mode

CHANGELOG.md

@@ -6,6 +6,7 @@ New Features
 - Adds ``eq_fixed`` argument to ``BoundaryError`` to remove the equilibrium from the optimization. This can be used instead of adding a ``FixParameter(eq)`` constraint.
 - Adds `check_intersection` argument to `initialize_modular_coils`, `initialize_helical_coils` and `initialize_saddle_coils`
 - Default value of `check_intersection` for coil related functions now defaults to False (no check). Previously, the default was True, and this was causing redundant checks.
+- Added warning for when ``deriv_mode="batched"`` is used in an ``ObjectiveFunction`` where one or more sub-objectives is using ``rev`` mode differentiation. Also adds more info about the derivative mode and Jacobian chunk sizes when building the objective with ``verbose>1``.
 
 Performance Improvements
 

desc/objectives/objective_funs.py

@@ -435,7 +435,7 @@ def _unjit(self):
                 pass
 
     @execute_on_cpu
-    def build(self, use_jit=None, verbose=1):
+    def build(self, use_jit=None, verbose=1):  # noqa: C901
         """Build the objective.
 
         Parameters
@@ -511,6 +511,21 @@ def build(self, use_jit=None, verbose=1):
             else:
                 self._deriv_mode = "blocked"
 
+        rev_objs = [
+            o.__class__.__name__ for o in self.objectives if o._deriv_mode == "rev"
+        ]
+        warnif(
+            len(rev_objs) > 0 and self._deriv_mode == "batched",
+            UserWarning,
+            "'batched' deriv_mode differentiates the whole ObjectiveFunction in "
+            "forward mode, but these sub-objectives are set to use reverse mode "
+            "(either automatically, from their input/output sizes, or by user): "
+            f"{rev_objs}. \n"
+            "In forward mode these may under-perform, or they may not "
+            "support forward mode. Consider 'blocked' deriv_mode. See the "
+            "sub-objective docstrings for details.",
+        )
+
         errorif(
             isposint(self._jac_chunk_size) and self._deriv_mode in ["blocked"],
             ValueError,
@@ -546,6 +561,12 @@ def build(self, use_jit=None, verbose=1):
 
         timer.stop("Objective build")
         if verbose > 1:
+            print(f"{self.name} deriv_mode : {self._deriv_mode}")
+            if self._deriv_mode == "batched":
+                print(f"{self.name} jac_chunk_size: {self._jac_chunk_size}")
+            else:
+                for o in self.objectives:
+                    print(f"{o.name} jac_chunk_size: {o._jac_chunk_size}")
             timer.disp("Objective build")
 
     def _set_things(self, things=None):

tests/test_objective_funs.py

@@ -2346,15 +2346,17 @@ def test_derivative_modes():
         jac_chunk_size="auto",
         use_jit=False,
     )
-    obj1.build()
+    with pytest.warns(UserWarning, match="batched"):
+        obj1.build()
     obj2.build()
     # check that default size works for blocked
     assert obj2.objectives[0]._jac_chunk_size == 2
     assert obj2.objectives[1]._jac_chunk_size is None
     assert obj2.objectives[2]._jac_chunk_size is None
     # hard to say what size auto will give, just check it is >0
     assert obj1._jac_chunk_size > 0
-    obj3.build()
+    with pytest.warns(UserWarning, match="batched"):
+        obj3.build()
     x = obj1.x(eq, surf)
     v = jnp.ones_like(x)
     g1 = obj1.grad(x)