custom autograd hooks #2987

groberts-flex · 2025-11-12T00:04:42Z

This adds custom user hooks for plugging into autograd. They are not exposed externally, but meant to be used/tested privately for now to help implement new/custom features and integrate with things like photonforge.

There are two hooks. The first is user_vjp which allows one to override the vjp calculation for a given structure that we already have in our library. For example, the sphere gradient could be computed with permittivity finite differences using this method. The second hook is numerical_structures which allows aa different custom integration. It allows the specification of a function that creates a tidy3d structure based on input parameters and then provides a vjp for those parameters. This is useful for cases where the adjoint fields are needed for that vjp. For example, you can implement a ring structure with a trimesh that depends on the inner and outer radius. This structure is inserted into the simulation and then the adjoint fields are available inside the vjp to write a numerical or other type of gradient calculation.

Greptile Overview

Greptile Summary

This PR introduces custom autograd hooks (user_vjp and numerical_structures) for gradient computation in electromagnetic simulations. The implementation adds two new internal functions run_custom and run_async_custom that wrap the existing public API.

Key Changes:

New user_vjp parameter allows overriding VJP calculations for existing structure geometries (e.g., computing sphere gradients via permittivity finite differences)
New numerical_structures parameter enables dynamic structure insertion with custom gradient functions, useful when adjoint fields are needed for gradient computation (e.g., trimesh-based ring structures)
Both hooks integrate into the existing adjoint gradient pipeline through modifications to _compute_derivatives and postprocess_adj
Public run and run_async functions remain unchanged; new functionality accessed via run_custom and run_async_custom

Issues Found:

Type checking using type() instead of isinstance() in two locations (lines 617, 630 of autograd.py)
Potential logic issue in structure.py where empty dict for vjp_fns may not correctly trigger custom VJP
Consider more explicit None checks vs truthiness checks for optional dict parameters

Testing:
Comprehensive numerical tests compare custom VJP/numerical structure gradients against finite differences for both single and batch runs.

Confidence Score: 3/5

Safe to merge after fixing type checking issues, though logic issue should be investigated
The PR adds significant new functionality with good test coverage, but contains syntax issues using type() instead of isinstance() that violate coding standards and could cause inheritance issues. There's also a potential logic bug in the custom VJP dispatch that needs verification. The implementation is complex but well-structured, and the hooks are appropriately kept internal for now.
Pay close attention to tidy3d/web/api/autograd/autograd.py (type checking violations) and tidy3d/components/structure.py (potential VJP dispatch logic issue)

Important Files Changed

File Analysis

Filename	Score	Overview
tidy3d/web/api/autograd/autograd.py	3/5	Core autograd entry point with new `run_custom` and `run_async_custom` functions. Contains type checking issues using `type()` instead of `isinstance()`. Complex flow for handling numerical structures and user VJP hooks.
tidy3d/web/api/autograd/backward.py	4/5	Implements backward pass for adjoint gradients. Adds support for numerical structures and user VJP hooks. New `updated_epsilon` function provides permittivity computation for finite differences.
tidy3d/web/api/autograd/types.py	5/5	New file defining type structures for user VJP and numerical structure metadata. Clean implementation using NamedTuples and dataclasses.
tidy3d/components/autograd/types.py	5/5	Adds `NumericalStructureInfo` dataclass and `CustomVJPPathType` to existing autograd types. Clean addition to type system.
tidy3d/components/structure.py	3/5	Modified `_compute_derivatives` to accept optional `vjp_fns` parameter. Has potential logic issue where empty dict for `vjp_fns` may not trigger custom VJP correctly.

Sequence Diagram

sequenceDiagram
    participant User
    participant run_custom
    participant setup_run
    participant _run_primitive
    participant run_webapi
    participant setup_adj
    participant postprocess_adj
    participant Structure
    
    User->>run_custom: simulation, numerical_structures, user_vjp
    
    alt has numerical_structures or traced fields
        run_custom->>setup_run: prepare simulation with hooks
        setup_run->>setup_run: insert numerical structures
        setup_run->>setup_run: extract traced fields
        setup_run-->>run_custom: sim_fields, prepared_sim, numerical_info
        
        run_custom->>_run_primitive: execute forward simulation
        _run_primitive->>run_webapi: run forward with aux_data
        run_webapi-->>_run_primitive: sim_data_fwd
        _run_primitive-->>run_custom: traced_fields_data
        
        Note over User,Structure: Backward Pass (when gradient needed)
        
        User->>setup_adj: request gradients with vjp seeds
        setup_adj->>setup_adj: create adjoint simulations
        setup_adj->>run_webapi: run adjoint simulations
        run_webapi-->>setup_adj: sim_data_adj
        setup_adj-->>User: adjoint data
        
        User->>postprocess_adj: compute VJPs from adjoint data
        
        alt user_vjp provided
            postprocess_adj->>Structure: _compute_derivatives(vjp_fns)
            Structure->>Structure: call user-provided VJP function
            Structure-->>postprocess_adj: custom gradients
        else numerical_structures
            postprocess_adj->>postprocess_adj: call numerical VJP function
            postprocess_adj->>postprocess_adj: compute finite difference
        else default
            postprocess_adj->>Structure: _compute_derivatives()
            Structure->>Structure: compute analytical gradients
            Structure-->>postprocess_adj: gradients
        end
        
        postprocess_adj-->>User: final VJP values
        
    else no traced fields or numerical structures
        run_custom->>run_webapi: regular simulation run
        run_webapi-->>User: simulation data
    end

Context used:

Rule from dashboard - Use isinstance() for type checking instead of type() to correctly handle inheritance. (source)

… and numerical_structures arguments to provide hooks into gradient computation for user-defined vjp calculation.

greptile-apps

_{22 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2025-11-12T00:08:04Z

tidy3d/web/api/autograd/autograd.py

            sim_dict[task_name] = sim
+
+        if user_vjp is not None:
+            if type(user_vjp) is not type(simulations):


syntax: Use isinstance() for type checking instead of type() to correctly handle inheritance

Suggested change

if type(user_vjp) is not type(simulations):

if not isinstance(user_vjp, type(simulations)):

Prompt To Fix With AI

This is a comment left during a code review. Path: tidy3d/web/api/autograd/autograd.py Line: 617:617 Comment: **syntax:** Use `isinstance()` for type checking instead of `type()` to correctly handle inheritance ```suggestion if not isinstance(user_vjp, type(simulations)): ``` How can I resolve this? If you propose a fix, please make it concise.

greptile-apps · 2025-11-12T00:08:04Z

tidy3d/web/api/autograd/autograd.py

+            user_vjp = user_vjp_dict
+
+        if numerical_structures is not None:
+            if type(numerical_structures) is not type(simulations):


syntax: Use isinstance() for type checking instead of type() to correctly handle inheritance

Suggested change

if type(numerical_structures) is not type(simulations):

if not isinstance(numerical_structures, type(simulations)):

Prompt To Fix With AI

This is a comment left during a code review. Path: tidy3d/web/api/autograd/autograd.py Line: 630:630 Comment: **syntax:** Use `isinstance()` for type checking instead of `type()` to correctly handle inheritance ```suggestion if not isinstance(numerical_structures, type(simulations)): ``` How can I resolve this? If you propose a fix, please make it concise.

greptile-apps · 2025-11-12T00:08:05Z

tidy3d/components/structure.py

+                    derivative_values_map.update(
+                        med_or_geo_field._compute_derivatives(derivative_info=info)


logic: Checking vjp_fns is not None before using it, but not checking if path_key in vjp_fns. If vjp_fns is an empty dict, the code will skip custom VJP even when one exists for the path_key

Consider: if vjp_fns and path_key in vjp_fns:

Prompt To Fix With AI

This is a comment left during a code review. Path: tidy3d/components/structure.py Line: 390:391 Comment: **logic:** Checking `vjp_fns is not None` before using it, but not checking if `path_key in vjp_fns`. If `vjp_fns` is an empty dict, the code will skip custom VJP even when one exists for the path_key Consider: `if vjp_fns and path_key in vjp_fns:` How can I resolve this? If you propose a fix, please make it concise.

github-actions · 2025-11-12T00:28:02Z

Diff Coverage

Diff: origin/develop...HEAD, staged and unstaged changes

tidy3d/components/autograd/derivative_utils.py (100%)
tidy3d/components/autograd/types.py (100%)
tidy3d/components/autograd/utils.py (50.0%): Missing lines 55-61
tidy3d/components/geometry/primitives.py (100%)
tidy3d/components/simulation.py (83.3%): Missing lines 4825,4829
tidy3d/components/structure.py (100%)
tidy3d/plugins/smatrix/run.py (96.4%): Missing lines 261
tidy3d/web/api/autograd/init.py (100%)
tidy3d/web/api/autograd/autograd.py (85.1%): Missing lines 73,75-78,80-81,87-88,92,94,105-106,108-111,113,120,123,154,167,175,375,445,618,631,910,1318,1323
tidy3d/web/api/autograd/backward.py (71.6%): Missing lines 124,153,267-268,270,280,285,355,392-394,398,417-419,423,427,429-430,433,437
tidy3d/web/api/autograd/constants.py (100%)
tidy3d/web/api/autograd/engine.py (100%)
tidy3d/web/api/autograd/types.py (100%)

Summary

Total: 376 lines
Missing: 61 lines
Coverage: 83%

tidy3d/components/autograd/utils.py

  51     if isinstance(value, Box):
  52         return True
  53     if isinstance(value, np.ndarray):
  54         return any(contains_tracer(v) for v in value.flat)
! 55     if isinstance(value, dict):
! 56         return any(contains_tracer(v) for v in value.values())
! 57     if isinstance(value, (list, tuple)):
! 58         return any(contains_tracer(v) for v in value)
! 59     if isinstance(value, Iterable) and not isinstance(value, (str, bytes)):
! 60         return any(contains_tracer(v) for v in value)
! 61     return False
  62 
  63 
  64 def pack_complex_vec(z):
  65     """Ravel [Re(z); Im(z)] into one real vector (autograd-safe)."""

tidy3d/components/simulation.py

  4821         numerical_indices = set()
  4822 
  4823         for namespace, index, *fields in sim_fields_keys:
  4824             if namespace not in {"structures", "numerical"}:
! 4825                 log.warning(
  4826                     "Encountered unknown namespace '%s' while creating adjoint monitors; ignoring.",
  4827                     namespace,
  4828                 )
! 4829                 continue
  4830 
  4831             if namespace == "structures":
  4832                 index_to_keys[index].append(fields)
  4833             elif namespace == "numerical":

tidy3d/plugins/smatrix/run.py

  257 
  258         return compose_modeler_data_from_batch_data(modeler=modeler, batch_data=sim_data_map)
  259 
  260     if numerical_structures is not None:
! 261         modeler = modeler.updated_copy(
  262             simulation=insert_numerical_structures_static(
  263                 simulation=modeler.simulation, numerical_structures=numerical_structures
  264             )
  265         )

tidy3d/web/api/autograd/autograd.py

  69     numerical_structures: dict[int, dict[str, typing.Any]],
  70 ) -> td.Simulation:
  71     """Return a Simulation with numerical structures inserted, without autograd metadata."""
  72 
! 73     structures = list(simulation.structures)
  74 
! 75     for index in sorted(numerical_structures):
! 76         config = numerical_structures[index]
! 77         func = config["function"]
! 78         params_input = config["parameters"]
  79 
! 80         try:
! 81             structure = func(get_static(params_input))
  82         except Exception as exc:  # pragma: no cover - defensive
  83             raise AdjointError(
  84                 f"Failed to construct numerical structure at index {index}: {exc}"
  85             ) from exc

  83             raise AdjointError(
  84                 f"Failed to construct numerical structure at index {index}: {exc}"
  85             ) from exc
  86 
! 87         if not isinstance(structure, td.Structure):
! 88             raise AdjointError(
  89                 "Numerical structure creation functions must return a tidy3d.Structure instance."
  90             )
  91 
! 92         structures.insert(index, structure)
  93 
! 94     return simulation.copy(update={"structures": structures})
  95 
  96 
  97 def _normalize_simulations_input(
  98     simulations: typing.Union[dict[str, td.Simulation], tuple[td.Simulation], list[td.Simulation]],

  101 
  102     if isinstance(simulations, dict):
  103         return simulations, {name: idx for idx, name in enumerate(simulations)}
  104 
! 105     normalized: dict[str, td.Simulation] = {}
! 106     name_mapping: dict[str, int] = {}
  107 
! 108     for idx, sim in enumerate(simulations):
! 109         task_name = Tidy3dStub(simulation=sim).get_default_task_name() + f"_{idx + 1}"
! 110         normalized[task_name] = sim
! 111         name_mapping[task_name] = idx
  112 
! 113     return normalized, name_mapping
  114 
  115 
  116 def normalize_user_vjp_spec(spec: tuple[CustomVJPPathType, ...]) -> typing.Optional[UserVjpSpec]:
  117     """Normalize a user-provided VJP specification into canonical tuple entries."""

  116 def normalize_user_vjp_spec(spec: tuple[CustomVJPPathType, ...]) -> typing.Optional[UserVjpSpec]:
  117     """Normalize a user-provided VJP specification into canonical tuple entries."""
  118 
  119     if spec is None:
! 120         return None
  121 
  122     if not spec:
! 123         return ()
  124 
  125     return tuple(UserVjpEntry(entry[0], (entry[1],), entry[2]) for entry in spec)
  126

  150     for cfg in numerical_structures.values():
  151         params = cfg.get("parameters")
  152         if contains_tracer(params):
  153             return True
! 154     return False
  155 
  156 
  157 def validate_numerical_structures(
  158     numerical_structures: dict[int, dict[str, typing.Any]],

  163 
  164     for index, numerical_config in numerical_structures.items():
  165         array_params = np.array(numerical_config["parameters"])
  166         if array_params.ndim != 1:
! 167             raise AdjointError(
  168                 f"Parameters for numerical structure index {index} must be 1D array-like."
  169             )
  170 
  171     # Reject user_vjp entries that try to target numerical namespace

  171     # Reject user_vjp entries that try to target numerical namespace
  172     if user_vjp:
  173         for entry in user_vjp:
  174             if entry.path and entry.path[0] == "numerical":
! 175                 raise AdjointError(
  176                     "Global 'user_vjp' cannot target 'numerical' namespace; specify VJP via numerical structure entry."
  177                 )
  178

  371         sim_dict = simulation.sim_dict
  372 
  373         numerical_structures_modeler = numerical_structures or {}
  374         if not numerical_structures_modeler and isinstance(numerical_structures_validated, dict):
! 375             numerical_structures_modeler = numerical_structures_validated
  376 
  377         should_use_component_autograd = any(
  378             is_valid_for_autograd(sim) for sim in sim_dict.values()
  379         ) or has_traced_numerical_structures(numerical_structures_modeler)

  441     simulation_static = simulation
  442     if isinstance(simulation, td.Simulation) and numerical_structures_validated:
  443         # if there are numerical_structures without traced parameters, we still want
  444         # to insert them into the simulation
! 445         simulation_static = insert_numerical_structures_static(
  446             simulation=simulation,
  447             numerical_structures=numerical_structures_validated,
  448         )

  614             sim_dict[task_name] = sim
  615 
  616         if user_vjp is not None:
  617             if type(user_vjp) is not type(simulations):
! 618                 raise AdjointError(
  619                     f"user_vjp type ({type(user_vjp)}) should match simulations type ({type(simulations)})"
  620                 )
  621 
  622             # set up the user_vjp_dict to have the same keys as the simulation dict

  627             user_vjp = user_vjp_dict
  628 
  629         if numerical_structures is not None:
  630             if type(numerical_structures) is not type(simulations):
! 631                 raise AdjointError(
  632                     f"numerical_structures type ({type(numerical_structures)}) should match simulations type ({type(simulations)})"
  633                 )
  634 
  635             # set up the numerical_structures_dict to have the same keys as the simulation dict

  906     for task_name in traced_fields_sim_dict.keys():
  907         traced_fields_data = traced_fields_data_dict[task_name]
  908         aux_data = aux_data_dict[task_name]
  909         if numerical_info_map.get(task_name) and AUX_KEY_NUMERICAL_STRUCTURES not in aux_data:
! 910             aux_data[AUX_KEY_NUMERICAL_STRUCTURES] = numerical_info_map[task_name]
  911         sim_data = postprocess_run(traced_fields_data=traced_fields_data, aux_data=aux_data)
  912         sim_data_dict[task_name] = sim_data
  913 
  914     return sim_data_dict

  1314 
  1315     task_names = data_fields_original_dict.keys()
  1316 
  1317     if numerical_structures_info is None:
! 1318         numerical_structures_info = {}
  1319 
  1320     if isinstance(user_vjp, dict):
  1321         user_vjp_map = user_vjp
  1322     else:
! 1323         user_vjp_map = dict.fromkeys(task_names, user_vjp)
  1324 
  1325     # get the fwd epsilon and field data from the cached aux_data
  1326     sim_data_orig_dict = {}
  1327     sim_data_fwd_dict = {}

tidy3d/web/api/autograd/backward.py

  120     user_vjp_lookup: dict[int, dict[typing.Hashable, typing.Callable[..., typing.Any]]] = {}
  121     if user_vjp:
  122         for structure_index, path, vjp_fn in user_vjp:
  123             if not path:
! 124                 continue
  125             field_key = path[0]
  126             user_vjp_lookup.setdefault(structure_index, {})[field_key] = vjp_fn
  127 
  128     # map of index into 'structures' and 'numerical' to the paths we need VJPs for

  149 
  150         if numerical_paths_raw:
  151             info = numerical_info.get(structure_index)
  152             if info is None:
! 153                 raise AdjointError(
  154                     f"Missing numerical structure metadata for index {structure_index}."
  155                 )
  156             numerical_vjp_fn = info.vjp
  157             numerical_params_static = tuple(get_static(param) for param in info.parameters)

  263         ) -> ScalarFieldDataArray:
  264             # Return the simulation permittivity for eps_box after replacing the geometry
  265             # for this structure with a new geometry. This is helpful for carrying out finite
  266             # difference permittivity computations
! 267             sim_orig = sim_data_orig.simulation
! 268             sim_orig_grid_spec = td.components.grid.grid_spec.GridSpec.from_grid(sim_orig.grid)
  269 
! 270             update_sim = sim_orig.updated_copy(
  271                 structures=[
  272                     sim_orig.structures[idx].updated_copy(geometry=replacement_geometry)
  273                     if idx == structure_index
  274                     else sim_orig.structures[idx]

  276                 ],
  277                 grid_spec=sim_orig_grid_spec,
  278             )
  279 
! 280             eps_by_f = [
  281                 update_sim.epsilon(box=eps_box, coord_key="centers", freq=f)
  282                 for f in adjoint_frequencies
  283             ]
  284 
! 285             return xr.concat(eps_by_f, dim="f").assign_coords(f=adjoint_frequencies)
  286 
  287         # get chunk size - if None, process all frequencies as one chunk
  288         freq_chunk_size = config.adjoint.solver_freq_chunk_size
  289         n_freqs = len(adjoint_frequencies)

  351                 select_adjoint_freqs: typing.Optional[FreqDataArray] = select_adjoint_freqs,
  352                 updated_epsilon_full: typing.Optional[typing.Callable] = updated_epsilon_full,
  353             ) -> ScalarFieldDataArray:
  354                 # Get permittivity function for a subset of frequencies
! 355                 return updated_epsilon_full(replacement_geometry).sel(f=select_adjoint_freqs)
  356 
  357             common_kwargs = {
  358                 "E_der_map": E_der_map_chunk,
  359                 "D_der_map": D_der_map_chunk,

  388                 vjp_chunk = structure._compute_derivatives(derivative_info_struct, vjp_fns=vjp_fns)
  389 
  390                 for path, value in vjp_chunk.items():
  391                     if path in vjp_value_map:
! 392                         existing = vjp_value_map[path]
! 393                         if isinstance(existing, (list, tuple)) and isinstance(value, (list, tuple)):
! 394                             vjp_value_map[path] = type(existing)(
  395                                 x + y for x, y in zip(existing, value)
  396                             )
  397                         else:
! 398                             vjp_value_map[path] = existing + value
  399                     else:
  400                         vjp_value_map[path] = value
  401 
  402             if numerical_paths_ordered and numerical_vjp_fn is not None:

  413                     gradient_items = (
  414                         (path, gradients.get(path)) for path in numerical_paths_ordered
  415                     )
  416                 else:
! 417                     gradients_seq = tuple(gradients)
! 418                     if len(gradients_seq) != len(numerical_paths_ordered):
! 419                         raise AdjointError(
  420                             f"User VJP for numerical structure index {structure_index} returned {len(gradients_seq)} gradients, "
  421                             f"expected {len(numerical_paths_ordered)}."
  422                         )
! 423                     gradient_items = zip(numerical_paths_ordered, gradients_seq)
  424 
  425                 for path, grad_value in gradient_items:
  426                     if grad_value is None:
! 427                         continue
  428                     if path in numerical_value_map:
! 429                         existing = numerical_value_map[path]
! 430                         if isinstance(existing, (list, tuple)) and isinstance(
  431                             grad_value, (list, tuple)
  432                         ):
! 433                             numerical_value_map[path] = type(existing)(
  434                                 x + y for x, y in zip(existing, grad_value)
  435                             )
  436                         else:
! 437                             numerical_value_map[path] = existing + grad_value
  438                     else:
  439                         numerical_value_map[path] = grad_value
  440 
  441         # store vjps in output map

marcorudolphflex

Really a cool feature!
Would work on readability/formatting.

And I wonder if some helpers to construct the user_vjp/numericals are feasible. When I look at the tests, it seems like one has to wire quite a lot around with stuff which is already pretty wired? Not sure how feasible, maybe also for a follow-up PR.
I know it should serve as an internal feature for now, but still I question if the usability could be easier.

marcorudolphflex · 2025-11-12T12:15:07Z

tidy3d/components/structure.py


-    def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap:
+    def _compute_derivatives(
+        self, derivative_info: DerivativeInfo, vjp_fns=None


missing type, consider add docstring update

marcorudolphflex · 2025-11-12T12:22:51Z

tidy3d/plugins/smatrix/run.py

 def _run_local(
    modeler: ComponentModelerType,
    path_dir: str = DEFAULT_DATA_DIR,
+    numerical_structures=None,


missing types, extend docstring

marcorudolphflex · 2025-11-12T12:27:19Z

tests/test_components/autograd/numerical/test_autograd_user_vjp.py

+    }
+    derivative_info_custom_medium = derivative_info.updated_copy(**update_kwargs)
+
+    def finite_difference_gradient(perturb_up, perturb_down, derivative_info_):


arguments unused - does not matter in this case but could be confusing after changes

marcorudolphflex · 2025-11-12T12:28:36Z

tidy3d/components/autograd/derivative_utils.py

    frequencies: ArrayLike
    """Frequencies at which the adjoint gradient should be computed."""

+    updated_epsilon: Callable


I'm leaning towards this being required since it's made for every structure that we call _compute_derivatives on

marcorudolphflex · 2025-11-12T12:33:56Z

tidy3d/components/geometry/primitives.py

    >>> b = Sphere(center=(1,2,3), radius=2)
    """

+    radius: TracedSize1D = pydantic.Field(


loses non-negativity constraint, consider adding validator

marcorudolphflex · 2025-11-12T12:47:51Z

tidy3d/web/api/autograd/autograd.py

    pay_type: typing.Union[PayType, str] = PayType.AUTO,
    priority: typing.Optional[int] = None,
    lazy: typing.Optional[bool] = None,
+    numerical_structures: typing.Optional[dict[int, dict[str, typing.Any]]] = None,


update docstring

marcorudolphflex · 2025-11-12T12:58:59Z

tidy3d/web/api/autograd/autograd.py

    pay_type: typing.Union[PayType, str] = PayType.AUTO,
    priority: typing.Optional[int] = None,
    lazy: typing.Optional[bool] = None,
+    numerical_structures: typing.Optional[


typing and handling is really hard to read/interpret... don't we want to go with class instances here (and throughout the code until some point)?

very good point - see above, but made these contain class instances in the new version

marcorudolphflex · 2025-11-12T13:04:37Z

tests/test_components/autograd/numerical/test_autograd_user_vjp.py

+        size=(dim_um, dim_um, thickness_um),
+    )
+
+    eval_fns, eval_fn_names = make_eval_fns()


marcorudolphflex · 2025-11-12T13:05:42Z

tests/test_components/autograd/numerical/test_autograd_user_vjp.py

+    """Test a variety of autograd permittivity gradients for DiffractionData by"""
+    """comparing them to numerical finite difference."""
+
+    test_number = test_parameters["test_number"]


marcorudolphflex · 2025-11-12T13:58:44Z

tests/test_components/autograd/numerical/test_autograd_numerical_structures.py

+        params_up[param_idx] += step_size
+        params_down[param_idx] -= step_size
+
+        rin_up = create_ring(params_up)


…r vjp interface and updated unit tests (not numerical tests yet)

…d unit tests

feature(adjoint): add custom internal run functions with new user_vjp…

30d1518

… and numerical_structures arguments to provide hooks into gradient computation for user-defined vjp calculation.

groberts-flex requested review from tylerflex and yaugenst-flex as code owners November 12, 2025 00:04

groberts-flex requested a review from marcorudolphflex November 12, 2025 00:05

greptile-apps bot reviewed Nov 12, 2025

View reviewed changes

marcorudolphflex requested changes Nov 12, 2025

View reviewed changes

Gregory Roberts and others added 5 commits November 13, 2025 15:56

part 1 of user_vjp interface overhaul; implemented round 1 of new use…

7e0756e

…r vjp interface and updated unit tests (not numerical tests yet)

update numerical tests with new user vjp

58af44d

new user vjp and numerical structures with updated numerical tests an…

165d4f9

…d unit tests

cleanup and test failures

c1ddcc7

docstrings, cleanup, review comments

9fe0414

	if type(user_vjp) is not type(simulations):
	if not isinstance(user_vjp, type(simulations)):

	if type(numerical_structures) is not type(simulations):
	if not isinstance(numerical_structures, type(simulations)):

		derivative_values_map.update(
		med_or_geo_field._compute_derivatives(derivative_info=info)

custom autograd hooks #2987

Are you sure you want to change the base?

custom autograd hooks #2987

Uh oh!

Conversation

groberts-flex commented Nov 12, 2025 • edited by greptile-apps bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 12, 2025

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 12, 2025

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Nov 12, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Nov 12, 2025

Diff Coverage

Diff: origin/develop...HEAD, staged and unstaged changes

Summary

tidy3d/components/autograd/utils.py

tidy3d/components/simulation.py

tidy3d/plugins/smatrix/run.py

tidy3d/web/api/autograd/autograd.py

tidy3d/web/api/autograd/backward.py

Uh oh!

marcorudolphflex left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

groberts-flex commented Nov 12, 2025 •

edited by greptile-apps bot

Loading