[1/n][Adjoint Module] Enable proper JAX support in adjoint solver by smartalecH · Pull Request #3280 · NanoComp/meep

smartalecH · 2026-08-25T19:00:32Z

This PR makes Jax a "first-class" citizen within Meep's adjoint model. While meep technically supported Jax already, the interface was buggy, lacked several important features, and didn't play nice with the existing autograd engine. This resolves that.

Allow users to specify their objective function using either autograd's numpy or jax's numpy. The backend handles the vjps automatically. You can't mix and match though (e.g. use meep's autograd filters with a jax objective).

from autograd import numpy as npa

def objective(mode_coeff, dft_fields):
    return npa.abs(mode_coeff)**2 + npa.sum(npa.abs(dft_fields)**2, axis=1)

OR

from jax import numpy as npj

def objective(mode_coeff, dft_fields):
    return npj.abs(mode_coeff)**2 + npj.sum(npj.abs(dft_fields)**2, axis=1)

Meep handles both just fine:

opt = mpa.OptimizationProblem(
    simulation=sim,
    objective_functions=[objective],
    objective_arguments=[mode_monitor, dft_monitor],
    design_regions=[design_region],
    frequencies=frequencies,
)

value, gradient = opt([rho])

Extend MeepJaxWrapper to properly support single-frequency objectives and vector-valued (multi-frequency) objectives. There was a bug in the single-frequency case. The multi-frequency case previously required the user to scalarize the objective (e.g. reduce across the frequency/wavelength dimension). You can now specify a vector-valued objective function (like you can with the standard OptimizationProblem API) and jax will return a dense jacobian:

def loss(rho, thickness, tilt):
    (dft,) = wrapped_meep([rho])
    return postprocess(dft, thickness, tilt)   # scalar

value, (d_rho, d_thickness, d_tilt) = jax.value_and_grad(
    loss, argnums=(0, 1, 2))(rho, 1.8, 8.0)

Package optimization parameters as pytrees (and return the gradient as an organized pytree). The end2end jax workflow allows you to chain operations before the meep simulations (preprocessing) and after (postprocessing). The user may want to collect gradients with parameters that tune these routines too. But bookeeping all of this is a pain. So we adopt the same approach e.g. fmmax uses via pytrees. You can pass a single pytree (a Dict/NamedTuple) to your objective function and the gradient routine will also produce a pytree. This also works with vector-valued objectives (e.g. for minimax optimization):

class Stack(NamedTuple):
    thickness: jnp.ndarray     # (3,) oxide, antireflection coating, air
    index: jnp.ndarray         # (3,)

class Params(NamedTuple):
    latent: jnp.ndarray        # (601,) design degrees of freedom
    beta: float                # preprocessing: projection strength
    stack: Stack               # post-processing: the layers above the chip
    tilt: float                # post-processing: fiber tilt, in degrees

def project(x, beta, eta=0.5):
    """The usual tanh projection, written in jax.numpy -- see the note below."""
    return (jnp.tanh(beta * eta) + jnp.tanh(beta * (x - eta))) / (
        jnp.tanh(beta * eta) + jnp.tanh(beta * (1 - eta))
    )

def loss(p):
    rho = project(conic_filter(p.latent), p.beta)   # preprocessing, in jax.numpy
    (ez, hx) = wrapped_meep([rho])
    field = propagate(ez, hx, p.stack)          # a stratified-media propagator, say
    return fiber_overlap(field, p.tilt)         # (nfreq,)

values, jacobian = mpa.value_and_jacobian(loss)(params)
jax.tree_util.tree_map(lambda leaf: leaf.shape, jacobian)

Various other bug fixes and removed antipatterns.

`grad` was imported in optimization_problem.py but never used. `_norm_fn`, `_reduce_fn`, and `_log_fn` in wrapper.py were never referenced. MeepJaxWrapper's `monitors` argument is annotated `List[EigenmodeCoefficient]`, but nothing in the wrapper is specific to that quantity -- it only calls the `ObjectiveQuantity` interface. Widen the annotation and the docstring to match what the code already does.

`OptimizationProblem.prepare_adjoint_run` called `autograd.jacobian` once per (objective function, objective argument) pair, and every `place_adjoint_source` implementation then contracted the frequency dimension of the result itself. But `autograd.jacobian` is one reverse pass per output component, so the old loop cost M*F reverse passes of the whole objective function -- for M objective arguments and F frequencies -- and materialized an (F, F, ...) array per argument only to sum it away. Meep runs one adjoint simulation per objective function and recovers a gradient at every frequency from it. That is only valid when component f of a frequency-vector-valued objective depends on the monitor values at frequency f alone -- the assumption already documented at the old contraction site in FourierFields. Under it, the array the adjoint sources need (the frequency diagonal of the Jacobian) is exactly the pullback of a cotangent of ones, and a single reverse pass over all arguments produces it. So: * `utils.objective_vjp` pulls a cotangent back through the objective function in one pass, returning the cotangents for every argument. * `OptimizationProblem._objective_cotangent` builds the seed, and now raises if an objective function returns a vector whose length is neither 1 nor nfreq -- previously that silently produced a misshapen contraction. * `place_adjoint_source` takes a cotangent shaped like the quantity's own value. The four open-coded contractions are gone; `_as_cotangent` validates the shape and still accepts the Jacobian shape with a DeprecationWarning. For M=4, F=5 this is 20 reverse passes down to 1, and peak memory for the objective gradient from O(F^2 N) to O(F N). Also unifies adjoint-source construction: `OptimizationProblem` now uses `utils.create_adjoint_sources` instead of its own inline loop, which is how it picks up the diagnostic for an all-zero gradient. That helper no longer forces cotangents through the field precision -- everything downstream of it works in double precision, so a single-precision build was round-tripping through complex64 for no reason.

…rote them `utils.objective_vjp` picks how to differentiate an objective function in three steps: a `vjp(cotangent, *args)` attribute on the function itself, if it has one; otherwise a backend registered with `register_vjp_backend` whose predicate recognizes the value the function returned; otherwise autograd. `wrapper` registers a backend for JAX on import, keyed on the return value being a `jax.Array`. The consequence is that an objective function written with `jax.numpy` is passed to `OptimizationProblem` exactly like an autograd one -- no wrapper type, no decorator, no annotation, and the two kinds can be mixed in one `objective_functions` list. The only difference between writing the two is which numpy is imported. Dispatching on the returned value rather than on the function keeps the choice precise: there is no way to tell that a plain Python function calls `jnp` internally without tracing it, and "try autograd, and if it throws assume JAX" would silently reinterpret real bugs in an autograd objective. `objective.py`, `optimization_problem.py`, and `utils.py` still have no knowledge of any framework but autograd, so JAX stays optional; a test asserts that nothing outside `wrapper.py` imports it. The explicit `vjp` attribute remains the escape hatch for an analytic derivative or a framework with no registered backend. Also warns when a double-precision Meep build is paired with JAX's default 32-bit mode, whose symptom is otherwise a finite-difference check that agrees to only three or four digits and looks like a physics bug.

…e monitors `utils.gather_monitor_values` stacked every monitor's value into one rank-2 (monitor, frequency) array and asserted the result was at most 2-D, so a `FourierFields` or `Near2FarFields` monitor -- which contributes a whole plane or a set of far points -- could not be used with `MeepJaxWrapper` at all. It now returns a tuple of per-monitor arrays when the shapes are heterogeneous, and keeps stacking when every monitor yields a single value per frequency, so the `monitor_values[i, :]` indexing in the module docstring and in user scripts is unaffected. `jax.custom_vjp` handles either structure, and the cotangent comes back matching it; `create_adjoint_sources` already iterates rather than indexing an array, so it needed no change. Drops `_make_at_least_nd`, whose only caller was the stacking path. The precision warning added with JaxObjective now also fires from MeepJaxWrapper, which has the same exposure.

test_adjoint_protocol.py checks the plumbing rather than the physics, and all but the last class runs without an FDTD simulation: * one VJP reproduces the contracted Jacobian for each of the value shapes the built-in objective quantities use; * the reverse pass runs once, not once per argument per frequency (asserted by counting VJP traversals through an identity primitive, since autograd.jacobian traces the forward pass once and replays the tape); * autograd and JAX agree on the complex cotangent convention for a real-valued objective of complex monitor values -- the case where a mismatch would silently yield a gradient of the right magnitude pointing the wrong way. Agreement is to 1e-12, against 1e+02 for the conjugate; * place_adjoint_source validates its argument's shape and still accepts the old Jacobian shape with one DeprecationWarning; * gather_monitor_values stacks homogeneous values and tuples heterogeneous ones; * no module of meep.adjoint outside wrapper.py imports jax, so the optional dependency cannot erode unnoticed; * end to end, a JaxObjective through OptimizationProblem matches the autograd version to 1e-10 and matches a finite difference. The end-to-end case uses two FourierFields monitors of different rank rather than a mode monitor: MPB warm-starts from the previous solve, so an otherwise identical repeat run reproduces a mode coefficient only to ~1e-8, which would put a floor under the framework-equivalence comparison.

Adds a section to the adjoint tutorial covering what an objective function receives and returns, why a frequency-vector-valued objective must be block diagonal in frequency for the single-adjoint-run formulation to hold, and how to use JaxObjective. Also spells out the distinction between JaxObjective, where Meep owns the optimization loop, and MeepJaxWrapper, where JAX does.

`DesignRegion.get_gradient` ends in `onp.squeeze(grad).T`, so with one frequency it returns a 1-D array of length num_design_params rather than (num_design_params, 1) -- the frequency axis is gone. `calculate_vjps` then contracted axis 1 and raised: AxisError: axis 1 is out of bounds for array of dimension 1 Every gradient through `MeepJaxWrapper` at a single frequency failed this way. `OptimizationProblem` is unaffected because it calls `get_gradient` directly and the squeeze is its documented return shape; only the wrapper goes through `calculate_vjps`. The existing wrapper tests parameterize over 3 and 4 frequencies, so nothing covered it. Restore the axis from the frequency count instead of relying on its survival. Verified against `OptimizationProblem` on the same scalar loss: gradients now agree to 6e-16 at one frequency, matching the 8e-16 already obtained at 3 and 5. Adds single-frequency cases to the wrapper's finite-difference test.

The JAX path could only differentiate a scalar loss, so it had no way to do worst-case (minimax) optimization over a bandwidth -- which needs a separate gradient per frequency, and which `OptimizationProblem` has always supported by returning a gradient with a frequency axis. `meep.adjoint.value_and_jacobian` is the counterpart of `jax.value_and_grad` for that case. The loss function is written exactly as it would be for a scalar objective: def loss(rho, thickness, tilt): (dft,) = wrapped_meep([rho]) return postprocess(dft, thickness, tilt) # (nfreq,) values, grads = mpa.value_and_jacobian(loss, argnums=(0, 1, 2))(rho, 1.8, 8.0) Every leaf of `grads` is the parameter's shape with a leading frequency axis. Parameters that never reach Meep -- a layer stack handed to a propagator, a fiber tilt -- are differentiated alongside the design weights, and pytrees work throughout. It costs one forward and one adjoint simulation regardless of the number of frequencies. `jax.jacrev` cannot achieve that: it evaluates the reverse pass once per output component, and here each evaluation is a full timestepping run. So this is a transform rather than something a JAX transformation does for itself. It works by splitting the loss at its `MeepJaxWrapper` call, which the wrapper cooperates with: the first pass runs the forward simulation and records what crossed the boundary, and the traced passes replay that recording instead of simulating. That leaves the design mapping before the call and the post-processing after it as ordinary JAX, so the explicit dependence on the parameters comes from `jax.jacrev` and the design rows are carried back with `jax.vmap` -- neither of which touches Meep. Only the frequency diagonal of the objective's Jacobian, obtained by seeding the reverse pass with ones, drives the one adjoint run. The block-diagonality requirement therefore binds only on the dependence of the loss on the monitor values, not on parameters that bypass the simulation. Uses the `sum_freq_partials=False` branch of `calculate_vjps`, which until now was unreachable. Verified against `OptimizationProblem`'s per-frequency gradient on the same objective: rows agree to 1.1e-15, while differing from each other by 29%, so the comparison is not vacuous. The tests compare each row against `jax.grad` of the corresponding component -- the same quantity by a route costing one adjoint run per row -- and assert the simulation counts are exactly one forward and one adjoint.

The pytree support was asserted in prose but only exercised with a flat dict. Two gaps: * no test had a parameter reaching the loss on *both* sides of the simulation -- through the design region and directly in post-processing -- so the sum of the adjoint and direct contributions, which value_and_jacobian's whole decomposition rests on, was never checked. Added, with a regularizer whose direct contribution is analytically known, plus a guard asserting that contribution is large enough here that dropping it would fail the test; * the tutorial described pytrees without showing one. Added a NamedTuple example with the tree of gradient shapes it returns, verified to be literally what the code prints. The pytree test now nests a NamedTuple inside a dict, matching the form users will actually write and the form the tutorial documents.

Formatting only, from the pinned pre-commit hook. Also adds a trailing comma after `**kwargs` in EigenmodeCoefficient.__init__: black picks its target version per file, and the f-strings added to objective.py in this branch move it to py36+, where that comma is legal.

… it is defined The section introduced `wrapped_meep` in a code example 88 lines before `MeepJaxWrapper` was described, and buried what that class actually is in a closing aside. Reordered to follow the way someone would come to it: 1. objective functions for OptimizationProblem, the interface that already existed; 2. the per-frequency gradient it returns, and what minimax needs from it, pointing at the mode_converter epigraph example; 3. that the objective may be written in jax.numpy instead, or carry its own pullback -- Meep does not care which framework produced it; 4. MeepJaxWrapper, what it inverts about the above, and the same optimizations through it; 5. value_and_jacobian for the broadband case; 6. the pytree parameterization, framed around the three groups it holds: preprocessing, the design degrees of freedom, and post-processing. Also fixes an example that could not have run: it called mpa.tanh_projection and mpa.conic_filter inside a JAX loss, but those are written with autograd.numpy and raise TracerArrayConversionError on a JAX tracer. The projection is now written in jax.numpy, with a note about the constraint. No test used the filters this way. The warning was first written as a `!!!` admonition, which mkdocs would have rendered as literal text -- the admonition extension is not enabled in mkdocs.yml, and no other page uses one.

The guard asserted that only wrapper.py may import JAX. That is stricter than the property worth protecting, which is that no module importing JAX is reachable unless the guarded import in __init__.py succeeded. Any new JAX-dependent module placed inside the same ModuleNotFoundError block satisfies that and would still have failed the old test, pushing whoever hit it toward either weakening the check or adding names to it by hand. It now parses the guarded imports out of __init__.py and asserts that the set of modules importing JAX is contained in them, so the check keeps working as the package grows and still fails if an unguarded module picks up a JAX import.

test_adjoint_protocol failed the MPI single-precision job. Two of its comparisons run the same problem twice -- once with an autograd objective and once with a JAX one -- and asserted agreement to 1e-12 and 1e-10. Those two routes reach the same gradient by different arithmetic, so they can only agree to whatever the monitor values themselves carry, which is about seven digits when Meep is built in single precision rather than the fifteen a double-precision build gives. The tolerances now branch on mp.is_single_precision(), as the finite-difference comparison in the same class already did. The equivalent comparison in test_adjoint_jax had the same hard 1e-13 and is branched too. It has not failed anywhere yet only because that file sits in ADJOINT_TESTS, which has no single-precision job; the latent problem is the same. Both remain three orders tighter than the finite-difference tolerances in the same tests, so they still fail if the two frameworks genuinely disagree.

oskooi

LGTM.

stevengj · 2026-08-28T20:11:27Z

I remember us having big problems with Jax and MPI on multi-core nodes when @ianwilliamson first implemented jax support. In particular, JAX tried to grab all of the CPUs on the node, which is bad if you are doing MPI intra-node parallelism. Is there a way to disable that these days?

smartalecH added 2 commits August 25, 2026 11:58

smartalecH changed the title ~~Differentiate adjoint objective functions with a VJP; support JAX objective functions~~ Enable proper JAX support in adjoint solver Aug 25, 2026

smartalecH added 4 commits August 25, 2026 12:53

smartalecH force-pushed the feat/adjoint-vjp-protocol branch from bdfbcc5 to e272cea Compare August 25, 2026 20:06

smartalecH added 5 commits August 25, 2026 13:18

smartalecH changed the title ~~Enable proper JAX support in adjoint solver~~ [1/n][Adjoint Module] Enable proper JAX support in adjoint solver Aug 26, 2026

smartalecH added 2 commits August 26, 2026 10:57

oskooi approved these changes Aug 28, 2026

View reviewed changes

Navigation Menu

Uh oh!

[1/n][Adjoint Module] Enable proper JAX support in adjoint solver - #3280

[1/n][Adjoint Module] Enable proper JAX support in adjoint solver#3280
smartalecH wants to merge 13 commits into
NanoComp:masterfrom
smartalecH:feat/adjoint-vjp-protocol

smartalecH commented Aug 25, 2026 •

edited

Loading

Uh oh!

oskooi left a comment

Uh oh!

stevengj commented Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

smartalecH commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

oskooi left a comment

Choose a reason for hiding this comment

Uh oh!

stevengj commented Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

smartalecH commented Aug 25, 2026 •

edited

Loading