maint[autograd]: warn instead of error if autograd tracers enter center/size of simulation, monitor, or source

Author: yaugenst-flex

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Changed
+- Supplying autograd-traced values to geometric fields (`center`, `size`) of simulations, monitors, and sources now logs a warning and falls back to the static value instead of erroring.
 - Attempting to differentiate server-side field projections now raises a clear error instead of silently failing.
 
 ## [2.8.3] - 2025-04-24

tests/test_components/test_autograd.py

@@ -2167,3 +2167,21 @@ def objective(args):
     # model is called without a frequency
     with AssertLogLevel("INFO"):
         grad = ag.grad(objective)(params0)
+
+
+def test_sim_traced_center_size(use_emulated_run):
+    fn_dict = get_functions(args[0][0], args[0][1])
+    make_sim = fn_dict["sim"]
+    postprocess = fn_dict["postprocess"]
+    base_sim = make_sim(params0)
+
+    def objective(center, size):
+        sim = base_sim.updated_copy(center=center, size=size)
+        sim_data = run_emulated(sim, task_name="adjoint_test")
+        return postprocess(sim_data)
+
+    with AssertLogLevel("WARNING", contains_str="autograd tracer"):
+        grad = ag.grad(objective, argnum=0)(base_sim.center, base_sim.size)
+
+    with AssertLogLevel("WARNING", contains_str="autograd tracer"):
+        grad = ag.grad(objective, argnum=1)(base_sim.center, base_sim.size)

tidy3d/components/base_sim/monitor.py

@@ -9,6 +9,7 @@
 from ..base import cached_property
 from ..geometry.base import Box
 from ..types import ArrayFloat1D, Axis, Numpy
+from ..validators import _warn_unsupported_traced_argument
 from ..viz import PlotParams, plot_params_monitor
 
 
@@ -22,6 +23,9 @@ class AbstractMonitor(Box, ABC):
         min_length=1,
     )
 
+    _warn_traced_center = _warn_unsupported_traced_argument("center")
+    _warn_traced_size = _warn_unsupported_traced_argument("size")
+
     @cached_property
     def plot_params(self) -> PlotParams:
         """Default parameters for plotting a Monitor object."""

tidy3d/components/base_sim/simulation.py

@@ -17,7 +17,11 @@
 from ..scene import Scene
 from ..structure import Structure
 from ..types import TYPE_TAG_STR, Ax, Axis, Bound, LengthUnit, Symmetry
-from ..validators import assert_objects_in_sim_bounds, assert_unique_names
+from ..validators import (
+    _warn_unsupported_traced_argument,
+    assert_objects_in_sim_bounds,
+    assert_unique_names,
+)
 from ..viz import (
     PlotParams,
     add_ax_if_none,
@@ -137,6 +141,9 @@ def _update_simulation(cls, values):
     _monitors_in_bounds = assert_objects_in_sim_bounds("monitors", strict_inequality=True)
     _structures_in_bounds = assert_objects_in_sim_bounds("structures", error=False)
 
+    _warn_traced_center = _warn_unsupported_traced_argument("center")
+    _warn_traced_size = _warn_unsupported_traced_argument("size")
+
     @pd.validator("structures", always=True)
     @skip_if_fields_missing(["size", "center"])
     def _structures_not_at_edges(cls, val, values):

tidy3d/components/source/base.py

@@ -11,7 +11,7 @@
 from ..base_sim.source import AbstractSource
 from ..geometry.base import Box
 from ..types import TYPE_TAG_STR, Ax
-from ..validators import _assert_min_freq
+from ..validators import _assert_min_freq, _warn_unsupported_traced_argument
 from ..viz import (
     ARROW_ALPHA,
     ARROW_COLOR_POLARIZATION,
@@ -58,6 +58,9 @@ def _pol_vector(self) -> Tuple[float, float, float]:
         """Returns a vector indicating the source polarization for arrow plotting, if not None."""
         return None
 
+    _warn_traced_center = _warn_unsupported_traced_argument("center")
+    _warn_traced_size = _warn_unsupported_traced_argument("size")
+
     @pydantic.validator("source_time", always=True)
     def _freqs_lower_bound(cls, val):
         """Raise validation error if central frequency is too low."""

tidy3d/components/validators.py

@@ -2,9 +2,11 @@
 
 import numpy as np
 import pydantic.v1 as pydantic
+from autograd.tracer import isbox
 
 from ..exceptions import SetupError, ValidationError
 from ..log import log
+from .autograd.utils import get_static
 from .base import DATA_ARRAY_MAP, skip_if_fields_missing
 from .data.dataset import Dataset, FieldDataset
 from .geometry.base import Box
@@ -464,3 +466,21 @@ def validate_mode_plane_radius(mode_spec: ModeSpec, plane: Box, msg_prefix: str
             f"{msg_prefix} bend radius is smaller than half the mode plane size "
             "along the radial axis, which can produce wrong results."
         )
+
+
+def _warn_unsupported_traced_argument(name: str):
+    @pydantic.validator(name, always=True, allow_reuse=True)
+    def _warn_traced_arg(cls, val, values):
+        if isbox(val):
+            log.warning(
+                f"Field '{name}' of '{cls.__name__}' received an autograd tracer "
+                f"(i.e., a value being tracked for automatic differentiation). "
+                f"Automatic differentiation through this field is unsupported, "
+                f"so the tracer has been converted to its static value. "
+                f"If you want to avoid this warning, you manually unbox the value "
+                f"using the 'autograd.tracer.getval' function before passing it to Tidy3D."
+            )
+            return get_static(val)
+        return val
+
+    return _warn_traced_arg