From f8a5efb05443b417044de5195cdb982097ab0099 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 6 Aug 2026 18:19:11 -0400 Subject: [PATCH 1/5] add warning for batched mode when a sub objective has rev mode --- desc/objectives/objective_funs.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 9ec99d6808..68fbf8c75b 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -511,6 +511,17 @@ def build(self, use_jit=None, verbose=1): else: self._deriv_mode = "blocked" + warnif( + any((obj._deriv_mode == "rev") for obj in self.objectives) + and self._deriv_mode == "batched", + UserWarning, + "Using batched derivative mode, which performs fwd mode " + "differentiation of the entire objective, but one or more of the" + ' sub-objectives have deriv_mode="rev". It is recommended to use ' + 'deriv_mode="blocked", as these objective performances may suffer ' + "in fwd mode.", + ) + errorif( isposint(self._jac_chunk_size) and self._deriv_mode in ["blocked"], ValueError, From f3830f8c148dfcb3a3568f3b7779c6dc7a5b7e4d Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 6 Aug 2026 18:58:22 -0400 Subject: [PATCH 2/5] update warning to include warning that sub-objective may not support rev mode --- desc/objectives/objective_funs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 68fbf8c75b..47257e7c3b 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -519,7 +519,7 @@ def build(self, use_jit=None, verbose=1): "differentiation of the entire objective, but one or more of the" ' sub-objectives have deriv_mode="rev". It is recommended to use ' 'deriv_mode="blocked", as these objective performances may suffer ' - "in fwd mode.", + "in fwd mode, or may not even support reverse mode at all.", ) errorif( From 61b868d62213bbf10584b3a646d959ac652412a4 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 10 Aug 2026 12:02:54 -0400 Subject: [PATCH 3/5] make message clearer, add printout of which objectives use rev mode, update changelog --- CHANGELOG.md | 1 + desc/objectives/objective_funs.py | 25 +++++++++++++++++-------- tests/test_objective_funs.py | 6 ++++-- 3 files changed, 22 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dfeee4bf11..7d5dfa2485 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 47257e7c3b..4ef7c3ab08 100644 --- a/desc/objectives/objective_funs.py +++ b/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,15 +511,18 @@ 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( - any((obj._deriv_mode == "rev") for obj in self.objectives) - and self._deriv_mode == "batched", + len(rev_objs) > 0 and self._deriv_mode == "batched", UserWarning, - "Using batched derivative mode, which performs fwd mode " - "differentiation of the entire objective, but one or more of the" - ' sub-objectives have deriv_mode="rev". It is recommended to use ' - 'deriv_mode="blocked", as these objective performances may suffer ' - "in fwd mode, or may not even support reverse mode at all.", + "'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}. In forward mode these may under-perform, or may not " + "work at all. Consider 'blocked' deriv_mode. See the sub-objective " + "docstrings for details.", ) errorif( @@ -557,6 +560,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): diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index f97bc93a27..4ab143d47b 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -2346,7 +2346,8 @@ 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 @@ -2354,7 +2355,8 @@ def test_derivative_modes(): 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) From 66b931872d8d7f529e33cbcc59ae8c6500532385 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 10 Aug 2026 13:05:34 -0400 Subject: [PATCH 4/5] tweak message --- desc/objectives/objective_funs.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 4ef7c3ab08..854cb83263 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -520,9 +520,10 @@ def build(self, use_jit=None, verbose=1): # noqa: C901 "'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}. In forward mode these may under-perform, or may not " - "work at all. Consider 'blocked' deriv_mode. See the sub-objective " - "docstrings for details.", + 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( From f163f9baadb387855034dffdcc94d23493cfaab2 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Thu, 13 Aug 2026 12:33:15 -0400 Subject: [PATCH 5/5] update changelog --- CHANGELOG.md | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e12640cd6c..dab6284414 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,15 +1,19 @@ Changelog ========= +New Features + +- 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 - Improves memory management to reduce the base memory used during optimization while using `lsq-exact`, `lsq-auglag` and `fmin-auglag` optimizers. - Bug Fixes - Fixes bug in ``auglag`` optimizers which prevented them from accepting solver hyperparameters. + v0.17.3 ------- @@ -18,7 +22,6 @@ 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