diff --git a/essos/coils.py b/essos/coils.py index 4351131..ec1cba1 100644 --- a/essos/coils.py +++ b/essos/coils.py @@ -633,6 +633,9 @@ def currents_scale(self, new_currents_scale): # dofs_currents property and setter @property def dofs_currents(self): + # Sentinel leaf during PyTree traversal: pass through, don't scale. + if self._dofs_currents_raw is None or isinstance(self._dofs_currents_raw, bool): + return self._dofs_currents_raw if self._dofs_currents is None: dofs_currents = self.dofs_currents_raw / self.currents_scale if not isinstance(dofs_currents, jax.core.Tracer):