fix[autograd]: ensure consistent task name mapping for batch autograd

Author: yaugenst-flex

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 - Gradient inaccuracy when a multi-frequency monitor is used but a single frequency is selected.
 - Revert single cell center approximation for custom medium gradient.
+- Potential task name mismatches between forward and adjoint simulations in batch tasks.
 
 ## [2.7.7] - 2024-11-15
 

tests/test_components/test_autograd.py

@@ -755,7 +755,7 @@ def test_autograd_async(use_emulated_run, structure_key, monitor_key):
     make_sim = fn_dict["sim"]
     postprocess = fn_dict["postprocess"]
 
-    task_names = {"1", "2", "3", "4"}
+    task_names = {"test_a", "adjoint", "task1", "_test"}
 
     def objective(*args):
         sims = {task_name: make_sim(*args) for task_name in task_names}

tidy3d/web/api/autograd/autograd.py

@@ -695,7 +695,7 @@ def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, Autograd
             return sim_fields_vjp_dict
 
         task_names_adj = list(sims_to_run.keys())
-        task_names_fwd = [name.rstrip("_adjoint") for name in task_names_adj]
+        task_names_fwd = [name.removesuffix("_adjoint") for name in task_names_adj]
 
         if local_gradient:
             # run adjoint simulation