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¶