Auto-deriving the chain rule with `request_AD`#

The previous tutorial hand-wrote the projectile’s first derivative — the actions closure lambda V: -mu * V encoding \(\partial\boldsymbol{a}/\partial\boldsymbol{v} = -\mu I\). That is the right call when the local Jacobian is simple. But when it isn’t — a constitutive law with many coupled terms, or a trained machine-learning surrogate we would never differentiate by hand — NEML2 can derive the chain rule for us by reverse-mode automatic differentiation.

In this tutorial, we’ll re-implement the same projectile model, but let request_AD supply the derivative: we write only the value forward, and the auto-derived chain rule comes out identical to the hand-written one.

Declaring `request_AD`#

A leaf opts in by calling self.request_AD() once at construction (in __post_init__), then writing a value-only forward — no v= branch, no actions. Everything else is unchanged: the same schema, the same _get_param read, the same physics. Here is the re-implemented projectile.py:

%%writefile projectile.py
"""Custom ``ProjectileAcceleration`` model — re-implemented with ``request_AD``.

Same model as the forward-operator tutorial (``a = g - mu * v``), but instead of
hand-writing the chain rule we declare ``request_AD`` and write only the value
``forward``; NEML2 supplies the first-order chain rule by reverse-mode autodiff.
"""

from __future__ import annotations

from neml2.factory import register_neml2_object
from neml2.models.model import Model
from neml2.schema import HitSchema, buffer, input, output, parameter
from neml2.types import Scalar, Vec


@register_neml2_object("ProjectileAcceleration")
class ProjectileAcceleration(Model):
    """Newton's second law for a projectile in a viscous medium, ``a = g - mu * v``,
    with the chain rule auto-derived by ``request_AD``.
    """

    hit = HitSchema(
        input("velocity", Vec, "Velocity of the projectile", attr="_v_name"),
        output("acceleration", Vec, "Acceleration of the projectile"),
        buffer(
            "gravitational_acceleration",
            Vec,
            "Gravity vector",
            attr="g",
            default=Vec.fill(0.0, -9.81, 0.0),
        ),
        parameter("dynamic_viscosity", Scalar, "Dynamic viscosity", attr="mu"),
    )

    _v_name: str
    g: Vec
    mu: Scalar

    def __post_init__(self):
        # Opt into auto-derived derivatives. With no arguments this covers every
        # (output, input) pair -- here the single pair d(acceleration)/d(velocity).
        # Pass outputs=[...] / inputs=[...] to auto-derive only a subset.
        self.request_AD()

    def forward(self, v_in: Vec, *promoted_params):  # type: ignore[override]
        # Read the drag coefficient through ``_get_param`` (promotion-compatible),
        # exactly as in the hand-written version.
        mu = self._get_param("mu", promoted_params, Scalar)

        # Value only: a = g - mu * v. No ``v=`` branch, no ``actions`` -- the
        # framework differentiates this forward by reverse-mode AD.
        return self.g - mu * v_in

Writing projectile.py

Two changes from the hand-written version:

__post_init__ calls self.request_AD(). With no arguments it covers every (output, input) pair — here \(\partial\boldsymbol{a}/\partial\boldsymbol{v}\). Pass outputs=[...] / inputs=[...] to auto-derive only a subset and hand-write the rest.
forward returns only the value. No v=None keyword, no apply_chain_rule. The framework intercepts the call, runs the value forward under reverse-mode autograd, and assembles the first-order chain rule from the result.

The parameter is still read through _get_param (promotion-compatible, exactly as before), and self.g is still a buffer read directly.

%%writefile input.i
[Models]
  [accel]
    type = ProjectileAcceleration
    velocity = 'v'
    acceleration = 'a'
    dynamic_viscosity = '0.001'
  []
[]

Writing input.i

Load and evaluate it exactly as we would any built-in type — the value forward behaves identically to the hand-written model:

import sys, os

sys.path.insert(0, os.getcwd())

import projectile  # registers ProjectileAcceleration with the native factory
import neml2

model = neml2.load_model("input.i", "accel")
model

ProjectileAcceleration()

import torch
from neml2.types import Vec

vel = Vec.fill(10.0, 2.0, 0.0)
accel = model(vel)
accel

Vec(data=tensor([-0.0100, -9.8120,  0.0000], dtype=torch.float64,
       grad_fn=<SubBackward0>), sub_batch_ndim=0, sub_batch_state=(), sub_batch_meta=(), k_ndim=0, k_state=(), k_pairing=())

The value is \(\boldsymbol{g} - \mu \boldsymbol{v} = (0, -9.81, 0) - 0.001 \cdot (10, 2, 0) = (-0.01, -9.812, 0)\) — identical to the hand-written model, as it must be.

First derivatives, for free#

We never wrote a v= branch, yet the model hands back directional derivatives just like the analytic one. Seed the identity on the velocity input and read off the full Jacobian column \(\partial\boldsymbol{a}/\partial\boldsymbol{v}\):

# The same identity-seed idiom as the hand-written tutorial — only now the
# Jacobian column is auto-derived rather than hand-coded.
seed = Vec(torch.eye(3, dtype=torch.float64))
accel, v_out = model(vel, v={"v": {"velocity_leaf": seed}})
v_out["a"]["velocity_leaf"].data

tensor([[-0.0010,  0.0000,  0.0000],
        [ 0.0000, -0.0010,  0.0000],
        [ 0.0000,  0.0000, -0.0010]], dtype=torch.float64)

This is exactly the \(-\mu I = -0.001\, I_3\) we hand-wrote last time — request_AD reproduced it from the value forward alone.

Verifying with `ModelUnitTest`#

The same ModelUnitTest harness pins the auto-derived model. It cross-checks every JVP against PyTorch’s autograd, so a positive JVP checks count is independent confirmation that the auto-derived chain rule is correct — not merely self-consistent:

%%writefile unit_test.i
[Tensors]
  [v_in]
    type = Python
    expr = 'Vec.fill(10.0, 2.0, 0.0)'
  []
  [a_expected]
    type = Python
    expr = 'Vec.fill(-0.01, -9.812, 0.0)'
  []
[]

[Models]
  [accel]
    type = ProjectileAcceleration
    velocity = 'v'
    acceleration = 'a'
    dynamic_viscosity = '0.001'
  []
[]

[Drivers]
  [unit]
    type = ModelUnitTest
    model = 'accel'
    input_Vec_names = 'v'
    input_Vec_values = 'v_in'
    output_Vec_names = 'a'
    output_Vec_values = 'a_expected'
  []
[]

Writing unit_test.i

from neml2.drivers.ModelUnitTest import ModelUnitTest

report = ModelUnitTest.from_file("unit_test.i").run()
print(f"value checks: {report.value_checks}, JVP checks: {report.jvp_checks}")

value checks: 1, JVP checks: 1

Where this really pays off: an ML surrogate#

The projectile’s \(-\mu I\) was easy to hand-write — request_AD only saved us a few lines. The real payoff is a forward operator whose Jacobian we would never want to derive by hand: a trained torch.nn.Module surrogate. The pattern is identical — declare request_AD, write only the value forward (now a network call):

class MySurrogate(Model):
    hit = HitSchema(...)            # same input/output declarations as usual

    def __post_init__(self):
        self.net = ...              # a trained torch.nn.Module
        self.request_AD()           # auto-derive d(output)/d(input) for all pairs

    def forward(self, x, *promoted_params):
        return type(x)(self.net(x.data))   # value only -- no `v=`, no `actions`

The auto-derived chain rule is indistinguishable from a hand-written one and behaves identically on every route — eager (py-eager / cpp-eager) and AOT-compiled (py-aoti / cpp-aoti / cpp-dispatch). It slots into the same forward-mode chain-rule graph the framework already uses: neighbouring analytic leaves keep their hand-written actions, and only the request_AD leaf’s reverse-mode local Jacobian is traced inline (and lowered through AOTInductor on the compiled routes). The tests/regression/_fixtures/SurrogateFlowRate.py fixture is a worked example wrapping a Python surrogate as a NEML2 flow rate.

A few things to know:

First-order only. request_AD supplies the v= channel (\(\partial\,\text{out}/\partial\,\text{in}\)). A leaf that must provide the second-order chain rule (i.e. one used inside a Normality wrap) still hand-writes it.
Reverse-mode under the hood. It is the one autodiff that survives torch.export → AOTInductor. If our differentiated path uses a saved-output op (exp / sqrt / tanh / reciprocal), route it through the AOTI-safe variants in neml2.types.functions (e.g. exp_ad) so the compiled routes lower (see the upstream-bug note in that module); eager is unaffected.
AOTI compile. Pass -d to neml2-compile to bake the derivative graph, exactly as for an analytic model — request_AD changes how the Jacobian is computed, not whether it is compiled. This holds even for a request_AD leaf inside an ImplicitUpdate residual: the Newton-step / implicit-function-theorem graphs differentiate the residual (which contains the leaf) and lower the same way.

Where to go next#

The next tutorial, Composing with existing models, shows how to glue several models together so a dependency resolver wires their inputs and outputs and threads the chain rule for us.

Auto-deriving the chain rule with request_AD#

Declaring request_AD#