From 78c699753bc41f46ea71a6a262e3aa2723a7c2df Mon Sep 17 00:00:00 2001 From: Tejas Date: Wed, 15 Jul 2026 02:30:59 -0500 Subject: [PATCH] Guard Coils.dofs_currents against None/bool sentinel leaves during PyTree traversal, fixing 'None is not a valid value for jnp.array' crash in tracing (from_simsopt and from_json paths) --- essos/coils.py | 3 +++ 1 file changed, 3 insertions(+) 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):