diff --git a/CHANGELOG.md b/CHANGELOG.md index 69a313b3a2..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 ------- diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 9ec99d6808..854cb83263 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,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): 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)