fix[adjoint]: fix tuple handling in autograd gradient calculations

Author: yaugenst-flex

CHANGELOG.md

@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Bug in `LayerRefinementSpec` that refines grids outside the layer region when one in-plane dimension is of size infinity.
 - Querying tasks was sometimes erroring unexpectedly.
 - Fixed automatic creation of missing output directories.
+- Bug in handling of tuple-type gradients that could lead to empty tuples or failing gradient calculations when differentiating w.r.t. (for instance) `td.Box.center`.
 
 ## [2.8.0] - 2025-03-04
 

tests/test_components/test_autograd.py

@@ -813,6 +813,83 @@ def objective(*args):
     assert anp.all(grad != 0.0), "some gradients are 0"
 
 
+class TestTupleGrads:
+    center0 = (0.0, 0.0, 0.0)
+    size0 = (0.5, 1.0, 1.5)
+
+    @staticmethod
+    def make_simulation(center: tuple, size: tuple) -> td.Simulation:
+        wavelength = 1.0
+        freq0 = td.C_0 / wavelength
+
+        src = td.PointDipole(
+            center=(-1.4, 0, 0),
+            source_time=td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10),
+            polarization="Ex",
+        )
+
+        mnt = td.FieldMonitor(
+            size=(0, 0, 1),
+            center=(1.4, 0, 0),
+            freqs=[freq0, freq0 + freq0 / 50],
+            name="fields",
+        )
+
+        scatterer = td.Structure(
+            geometry=td.Box(center=center, size=size),
+            medium=td.Medium(permittivity=3.0),
+        )
+
+        return td.Simulation(
+            size=(3, 3, 3),
+            run_time=2e-13,
+            structures=[scatterer],
+            sources=[src],
+            monitors=[mnt],
+            boundary_spec=td.BoundarySpec.all_sides(td.PML()),
+            grid_spec=td.GridSpec.auto(min_steps_per_wvl=30),
+        )
+
+    @pytest.mark.parametrize("run_async", [False, True])
+    @pytest.mark.parametrize("zero", [False, True])
+    @pytest.mark.parametrize("local_gradient", [False, True])
+    def test_zero_grad_tuple(self, use_emulated_run, run_async, zero, local_gradient, tmp_path):
+        """Checks that tuple gradients don't return empty tuples"""
+
+        def obj(center: tuple, size: tuple) -> float:
+            sim = self.make_simulation(center=center, size=size)
+            if run_async:
+                batch_data = web.run_async(
+                    {"lossy_test_async": sim},
+                    path_dir=tmp_path,
+                    local_gradient=local_gradient,
+                )
+                sim_data = list(batch_data.values())[0]
+            else:
+                sim_data = web.run(
+                    sim,
+                    task_name="lossy_test",
+                    local_gradient=local_gradient,
+                )
+            objval = anp.mean(sim_data["fields"].intensity.data).item()
+            if zero:
+                objval *= 0
+            return objval
+
+        d_power = ag.value_and_grad(obj, argnum=(0, 1))
+        val, (dp_dcenter, dp_dsize) = d_power(self.center0, self.size0)
+
+        assert len(dp_dcenter) == 3
+        assert len(dp_dsize) == 3
+
+        if zero:
+            assert np.allclose(dp_dcenter, 0)
+            assert np.allclose(dp_dsize, 0)
+        else:
+            assert not np.allclose(dp_dcenter, 0)
+            assert not np.allclose(dp_dsize, 0)
+
+
 @pytest.mark.parametrize("structure_key, monitor_key", args)
 def test_autograd_async_some_zero_grad(use_emulated_run, structure_key, monitor_key):
     """Test objective where only some simulations in batch have adjoint sources."""

tidy3d/web/api/autograd/autograd.py

@@ -648,7 +648,10 @@ def vjp(data_fields_vjp: AutogradFieldMap) -> AutogradFieldMap:
                 "simulation's output. If this is unexpected, please review your "
                 "setup or contact customer support for assistance."
             )
-            return {k: 0 * v for k, v in sim_fields_original.items()}
+            return {
+                k: (type(v)(0 * x for x in v) if isinstance(v, (list, tuple)) else 0 * v)
+                for k, v in sim_fields_original.items()
+            }
 
         # Run adjoint simulations in batch
         task_names_adj = [f"{task_name}_adjoint_{i}" for i in range(len(sims_adj))]
@@ -670,17 +673,16 @@ def vjp(data_fields_vjp: AutogradFieldMap) -> AutogradFieldMap:
             )
             td.log.info("Completed local batch adjoint simulations")
 
-            # sum partial derivatives from each adjoint simulation
+            # Process results from local gradient computation
+            vjp_fields_dict = {}
             for task_name_adj, sim_data_adj in batch_data_adj.items():
                 td.log.info(f"Processing VJP contribution from {task_name_adj}")
-                vjp_fields = postprocess_adj(
+                vjp_fields_dict[task_name_adj] = postprocess_adj(
                     sim_data_adj=sim_data_adj,
                     sim_data_orig=sim_data_orig,
                     sim_data_fwd=sim_data_fwd,
                     sim_fields_keys=sim_fields_keys,
                 )
-                for k, v in vjp_fields.items():
-                    vjp_traced_fields[k] = vjp_traced_fields.get(k, 0) + v
         else:
             td.log.info("Starting server-side batch of adjoint simulations ...")
 
@@ -699,15 +701,24 @@ def vjp(data_fields_vjp: AutogradFieldMap) -> AutogradFieldMap:
                 tname_adj: sim.updated_copy(simulation_type="autograd_bwd", deep=False)
                 for tname_adj, sim in sims_adj_dict.items()
             }
-            vjp_traced_fields_dict = _run_async_tidy3d_bwd(
+            vjp_fields_dict = _run_async_tidy3d_bwd(
                 simulations=sims_adj_dict,
                 **run_kwargs,
             )
             td.log.info("Completed server-side batch of adjoint simulations.")
 
-            for fields in vjp_traced_fields_dict.values():
-                for k, v in fields.items():
-                    vjp_traced_fields[k] = vjp_traced_fields.get(k, 0) + v
+        # Accumulate gradients from all adjoint simulations
+        for task_name_adj, vjp_fields in vjp_fields_dict.items():
+            td.log.info(f"Processing VJP contribution from {task_name_adj}")
+            for k, v in vjp_fields.items():
+                if k in vjp_traced_fields:
+                    val = vjp_traced_fields[k]
+                    if isinstance(val, (list, tuple)) and isinstance(v, (list, tuple)):
+                        vjp_traced_fields[k] = type(val)(x + y for x, y in zip(val, v))
+                    else:
+                        vjp_traced_fields[k] += v
+                else:
+                    vjp_traced_fields[k] = v
 
         td.log.debug(f"Computed gradients for {len(vjp_traced_fields)} fields")
         return vjp_traced_fields
@@ -765,7 +776,8 @@ def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, Autograd
             if not sims_adj:
                 td.log.debug(f"Adjoint simulation for task '{task_name}' contains no sources.")
                 sim_fields_vjp_dict[task_name] = {
-                    k: 0 * v for k, v in sim_fields_original_dict[task_name].items()
+                    k: (type(v)(0 * x for x in v) if isinstance(v, (list, tuple)) else 0 * v)
+                    for k, v in sim_fields_original_dict[task_name].items()
                 }
                 continue
 
@@ -781,6 +793,9 @@ def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, Autograd
             )
             return sim_fields_vjp_dict
 
+        # Dictionary to store VJP results from all adjoint simulations
+        vjp_results = {}
+
         if local_gradient:
             # Run all adjoint simulations in a single batch
             path_dir = Path(run_async_kwargs.pop("path_dir"))
@@ -791,28 +806,20 @@ def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, Autograd
                 all_sims_adj, path_dir=str(path_dir_adj), **run_async_kwargs
             )
 
-            # Process results for each original task
+            # Process results for each adjoint task
             for adj_task_name, sim_data_adj in batch_data_adj.items():
                 task_name = task_name_mapping[adj_task_name]
                 sim_data_orig = sim_data_orig_dict[task_name]
                 sim_data_fwd = sim_data_fwd_dict[task_name]
                 sim_fields_keys = sim_fields_keys_dict[task_name]
 
                 # Compute VJP contribution
-                sim_fields_vjp = postprocess_adj(
+                vjp_results[adj_task_name] = postprocess_adj(
                     sim_data_adj=sim_data_adj,
                     sim_data_orig=sim_data_orig,
                     sim_data_fwd=sim_data_fwd,
                     sim_fields_keys=sim_fields_keys,
                 )
-
-                # Sum contributions for each original task
-                if task_name in sim_fields_vjp_dict:
-                    for k, v in sim_fields_vjp.items():
-                        sim_fields_vjp_dict[task_name][k] += v
-                else:
-                    sim_fields_vjp_dict[task_name] = sim_fields_vjp
-
         else:
             # Set up parent tasks mapping for all adjoint simulations
             parent_tasks = {}
@@ -830,19 +837,27 @@ def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, Autograd
             }
 
             # Run all adjoint simulations in a single batch
-            sim_fields_vjp_dict_adj = _run_async_tidy3d_bwd(
+            vjp_results = _run_async_tidy3d_bwd(
                 simulations=all_sims_adj,
                 **run_async_kwargs,
             )
 
-            # Combine results for each original task
-            for adj_task_name, fields in sim_fields_vjp_dict_adj.items():
-                task_name = task_name_mapping[adj_task_name]
-                if task_name in sim_fields_vjp_dict:
-                    for k, v in fields.items():
+        # Accumulate gradients from all adjoint simulations
+        for adj_task_name, vjp_fields in vjp_results.items():
+            task_name = task_name_mapping[adj_task_name]
+
+            if task_name not in sim_fields_vjp_dict:
+                sim_fields_vjp_dict[task_name] = {}
+
+            for k, v in vjp_fields.items():
+                if k in sim_fields_vjp_dict[task_name]:
+                    val = sim_fields_vjp_dict[task_name][k]
+                    if isinstance(val, (list, tuple)) and isinstance(v, (list, tuple)):
+                        sim_fields_vjp_dict[task_name][k] = type(val)(x + y for x, y in zip(val, v))
+                    else:
                         sim_fields_vjp_dict[task_name][k] += v
                 else:
-                    sim_fields_vjp_dict[task_name] = fields
+                    sim_fields_vjp_dict[task_name][k] = v
 
         return sim_fields_vjp_dict