onnx-ad

Automatic differentiation of ONNX graphs, by source-code transformation: given a model, produce new ONNX models that compute its Jacobian-vector and vector-Jacobian products. Pure Python over the onnx protobuf, no runtime dependency, no framework in the loop.

python -m pip install onnx-ad

import onnx
from onnx_ad import forward, reverse, family

model = onnx.load("f.onnx")                      # inputs x -> outputs y
onnx.save(forward(model), "fwd_f.onnx")          # + fwd_x -> + fwd_y = J . fwd_x
onnx.save(reverse(model), "adj_f.onnx")          # + adj_y -> + adj_x = J^T . adj_y

family(model, "generated/f.onnx")                # the whole set CasADi discovers

A derivative model keeps the original signature as a prefix and appends the seeds, so the primal outputs stay available. Any number of seed directions is evaluated in a single pass.

Why differentiate ONNX

The third route is the one with no ceiling. A Jacobian-vector-product graph is itself an ONNX model, so it can be differentiated again, and consumers need no new capability.

Composition: second derivatives for free

The passes keep the primal graph and emit only ordinary ONNX operations, so their own output is differentiable. Forward-over-adjoint — the exact-Hessian building block — is the two passes composed, with no special casing:

adjoint = reverse(model)                 # x, adj_y -> adj_x
hessian = forward(adjoint)               # + fwd_x, fwd_adj_y -> + fwd_adj_x

That only works if the rule set is closed under its own output — the reverse of Gather is a ScatterND and the reverse of Conv a ConvTranspose, so both have rules too.

Repeated differentiation names itself the way CasADi's diff_prefix does: a model that already carries fwd_x gets fwd2_/nfwd2 next, so forward(forward(model)) needs no arguments.

CasADi

The conventions are CasADi's, by default, so an emitted family drops into its ONNX backend with nothing to adapt:

• names — fwd_<x>, adj_<y>, then fwd2_, adj2_, from the same rule FunctionInternal::diff_prefix applies, and seed dimensions nfwd, nadj, nfwd2;
• layout — CasADi reads an ONNX tensor as a matrix (rank 0/1 as a column, rank 2 directly, higher ranks flattened) and wants the seeds of an r-by-c value as one r-by-(nseed*c) matrix. That is what layout="casadi" emits. Pass layout="onnx" for the internal form instead, where the seed count is a trailing axis on the primal's own shape;
• files — family writes f.onnx, adj_f.onnx and fwd_adj_f.onnx, the <kind>_<filename> siblings the backend looks for beside a model.

f = casadi.GraphBuilder("generated/f.onnx").create("f")
f.reverse(1)(x, f(x), w)                            # from adj_f.onnx
casadi.hessian(casadi.dot(f(v), w), v)              # from fwd_adj_f.onnx

examples/torch_to_casadi.py exports a PyTorch model's primal only and generates the rest here; examples/casadi_side.py consumes it and checks gradients, Jacobians and exact Hessians against PyTorch. Verified against a CasADi build with WITH_ONNX=ON and WITH_ONNX_RUNTIME=ON, with CASADI_ONNXRUNTIME_LIB pointing at libonnxruntime.so.

Two traps when the primal comes from torch.onnx.export: pass external_data=False, or the weights land in a sidecar f.onnx.data that CasADi cannot follow (it hands the model to ONNX Runtime as bytes); and CasADi needs the complete forward signature of the adjoint, which is why family seeds adj_y as well as x.

Deliberately not offered: a jacobian pass. CasADi builds dense Jacobians from the adjoint itself, and a jac_ sibling would only duplicate that.

What it emits

Two walks over one rule table. Forward carries a tangent per value in graph order; reverse walks backwards, accumulating a contribution per value and summing where a value has several consumers. A value with no derivative is absent rather than zero, which is what keeps the emitted graph the size of the primal one — every weight in a network is such a value, and costs nothing.

Nonlinear rules read the primal tensors they need straight from the primal graph rather than recomputing them: the tangent of Tanh is (1 - y*y) * t, with y the tensor the primal Tanh already produced. The reverse model therefore stays a plain function of (x, adj_y) — no "uses output" convention, nothing to wire up.

Reverse mode's sharp edge is broadcasting: a contribution arrives shaped like the result and must be summed back over the axes the operand was broadcast along. Where the operand's shape is declared this is a static axis list; where it is symbolic the axes are computed at run time, so a dynamic batch dimension survives.

Operations with rules

Arithmetic Add Sub Mul Div Neg Pow Sum Mean Identity

Elementwise Exp Log Sqrt Reciprocal Abs Sign Sin Cos Tan Sinh Cosh Asin Acos Atan Asinh Acosh Atanh Erf Tanh Sigmoid Relu LeakyRelu Elu Selu Celu PRelu ThresholdedRelu Softplus Softsign Shrink HardSigmoid HardSwish Mish Gelu (exact and tanh)

Linear algebra MatMul Gemm Conv (strided, dilated, grouped, 1-D and up; weight and bias too) ConvTranspose

Shape Reshape Flatten Transpose Squeeze Unsqueeze Expand Concat Split Slice Pad Tile Gather GatherND ScatterND CumSum

Control flow If Scan Loop — see below

Reductions ReduceSum ReduceMean ReduceMax ReduceMin ReduceProd ReduceLogSumExp ReduceL1 ReduceL2 ReduceSumSquare

Selection Where Clip Min Max

Networks Softmax LogSoftmax LayerNormalization BatchNormalization (inference) Dropout (inference) Cast CastLike

Zero derivative, and allowed to consume differentiated values Shape Size NonZero Equal Greater Less GreaterOrEqual LessOrEqual And Or Xor Not the bitwise family ArgMax ArgMin IsNaN IsInf Floor Ceil Round Hardmax OneHot Det

An operation without a rule is an error only when a differentiated value reaches it — a Resize on a constant branch is fine.

Still missing: pooling (MaxPool, AveragePool, GlobalAveragePool), Einsum, Resize, ScatterElements, LpNormalization, InstanceNormalization and GroupNormalization. Out of scope: sparsity and training-mode operations.

Control flow

If, Scan and Loop are differentiated in place: the primal node is replaced by one whose subgraph also computes the derivative, rather than a derivative being built beside it.

• If — only the taken branch computes its derivative. An If on a constant condition is inlined before either pass runs.
• Scan — forward mode is one Scan carrying the tangent as extra state: two times the body, any number of seeds, nothing stored. Reverse mode tapes each iteration's input state and appends a second Scan that reads everything backwards, recomputes the body and runs its adjoint. A captured weight accumulates in that scan's state, so its adjoint costs O(|w|), not O(T·|w|).
• Loop — the trip count may depend on the data, which is only a problem until the primal has run: then it is the length of the tape. The reverse sweep is therefore Scan's, reading the tape backwards, behind an If for the loop that ran zero times.

Nesting works in both directions, and so does composition: forward(reverse(model)) through any of the three. The design, and what ONNX Runtime was checked to allow, is in CONTROL-FLOW.md; the slides draw it.

unroll(model) flattens every Scan and Loop whose trip count is known before the model runs. That is useful on its own for short, fixed-length recurrences — at the price of a graph that grows with the trip count — and it is what the control-flow rules are tested against: the Jacobian of the unrolled graph comes from entirely different code.

Against PyTorch

Every one of these exports (torch.onnx.export, dynamo=True, opset 18) differentiates in both modes, the resulting Jacobian matches torch.autograd.functional.jacobian to float32 precision, and family builds its second-order file:

Running the result

Reverse mode emits Transpose feeding MatMul, which ONNX Runtime's extended optimizer fuses into com.microsoft.FusedMatMul — a kernel registered for float only. On a double-precision model, load with the fusions off:

options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
session = ort.InferenceSession("adj_f.onnx", options)

Command line

onnx-ad forward f.onnx fwd_f.onnx
onnx-ad reverse f.onnx adj_f.onnx --inputs x --outputs y
onnx-ad family  f.onnx generated/f.onnx

Testing

Every rule is executed through ONNX Runtime and compared against references that know nothing of the rule table: central finite differences of the primal model, an analytic Jacobian written in numpy, and forward against reverse — J and J^T come from separate walks over separate rules, so their agreement to machine precision is a real check.

No PyTorch is involved in the unit tests; comparisons against it belong in an integration suite, not here.

python -m pip install -e ".[test]"
python -m unittest discover -s tests -v

License

MIT. Releasing is documented in RELEASE.md: PyPI Trusted Publishing, so no API token exists, and every artifact carries a PEP 740 provenance attestation.