Automatic Differentiation from Scratch: How PyTorch Computes Gradients in Physics-Informed Neural Networks

# Automatic Differentiation from Scratch: How Pytorch Computes Gradients in Physics-Informed Neural Networks
Source: [https://arxiv.org/html/2607.13042](https://arxiv.org/html/2607.13042)
Abdeladhim Tahimi CECA, Universidade Federal de Alagoas \(UFAL\), Rio Largo, Alagoas, 57100\-000, Brazil abdeladhim\.tahimi@ceca\.ufal\.br

###### Abstract

This paper traces, with explicit numerical values, how PyTorch’s automatic differentiation \(AD\) engine computes gradients for Physics\-Informed Neural Network \(PINN\) training — a setting that requires two levels of differentiation: computing the physics derivativey^′​\(t\)=d​y^/d​t\\hat\{y\}^\{\\prime\}\(t\)=d\\hat\{y\}/dtthrough the network, and computing parameter gradients∇θL\\nabla\_\{\\theta\}Lof a loss that itself depends ony^′​\(t\)\\hat\{y\}^\{\\prime\}\(t\)\. Using a 1\-3\-3\-1 multilayer perceptron and the initial value problemy′​\(t\)\+y​\(t\)=0y^\{\\prime\}\(t\)\+y\(t\)=0,y​\(0\)=1y\(0\)=1from\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\], we trace the complete pipeline at every node: the computational graph built during the forward pass, the reverse\-mode backward traversal that computes all 22 parameter gradients in a single pass, and the graph\-on\-graph mechanism by whichcreate\_graph=Trueenables correct differentiation through the physics\-informed residual\. Every adjoint value is verified against the hand derivations of\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\], connecting theP/QP/Qsensitivity framework to the vector–Jacobian products used by PyTorch’s autograd engine\.

Keywords:automatic differentiation, reverse\-mode automatic differentiation, backpropagation, computational graph, vector–Jacobian product, Physics\-Informed Neural Networks, PyTorch

MSC codes:65D25, 65\-01, 68T07, 65L05

## 1Introduction

A typical Physics\-Informed Neural Network \(PINN\) training loop fits in fewer than twenty lines of PyTorch\. The forward pass computes the network outputy^​\(t;θ\)\\hat\{y\}\(t;\\theta\); a call to[torch\.autograd\.grad](https://arxiv.org/html/2607.13042v1/torch.autograd.grad)withcreate\_graph=Truecomputes the time derivativey^′​\(t\)\\hat\{y\}^\{\\prime\}\(t\); the ODE residualR=y^′\+y^R=\\hat\{y\}^\{\\prime\}\+\\hat\{y\}is squared to form the loss; and[loss\.backward\(\)](https://arxiv.org/html/2607.13042v1/loss.backward())produces all parameter gradients in a single call\. The entire pipeline relies on automatic differentiation \(AD\) — yet what does each of these calls actually compute, and why is each flag necessary? Answering these questions requires opening the black box\.

##### Two differentiation tasks\.

Training a PINN requires two fundamentally different derivatives\. Consider the first\-order ODEy′​\(t\)\+y​\(t\)=0y^\{\\prime\}\(t\)\+y\(t\)=0with initial conditiony​\(0\)=1y\(0\)=1\. The PINN loss at a collocation pointtct\_\{c\}is

L=\(y^​\(tc;θ\)\+∂y^∂t\|tc\)2⏟LR:ODE residual\+λ​\(y^​\(0;θ\)−1\)2⏟LI​C:initial condition\.L=\\underbrace\{\\left\(\\hat\{y\}\(t\_\{c\};\\theta\)\+\\frac\{\\partial\\hat\{y\}\}\{\\partial t\}\\bigg\|\_\{t\_\{c\}\}\\right\)^\{2\}\}\_\{L\_\{R\}:\\;\\text\{ODE residual\}\}\+\\;\\lambda\\,\\underbrace\{\\bigl\(\\hat\{y\}\(0;\\theta\)\-1\\bigr\)^\{2\}\}\_\{L\_\{IC\}:\\;\\text\{initial condition\}\}\.\(1\)The loss involves the*physics derivative*∂y^/∂t\\partial\\hat\{y\}/\\partial t— how the network output changes when the inputttvaries, with all parameters held fixed — and training requires a second derivative, the*training gradient*∇θL\\nabla\_\{\\theta\}L, which describes how the loss changes when a parameterθk\\theta\_\{k\}is adjusted\. We writey^′\\hat\{y\}^\{\\prime\}ford​y^/d​td\\hat\{y\}/dt; since the parametersθ\\thetado not depend ontt, this coincides with∂y^/∂t\\partial\\hat\{y\}/\\partial t\.

In standard supervised learning, the loss depends only ony^\\hat\{y\}, and a single backward pass computes∇θL\\nabla\_\{\\theta\}L\. Here the loss depends on bothy^\\hat\{y\}andy^′\\hat\{y\}^\{\\prime\}, so the training gradient of the residual loss is

∂LR∂θk=2​R​\(∂y^∂θk\+∂y^′∂θk\)\.\\frac\{\\partial L\_\{R\}\}\{\\partial\\theta\_\{k\}\}=2R\\left\(\\frac\{\\partial\\hat\{y\}\}\{\\partial\\theta\_\{k\}\}\+\\frac\{\\partial\\hat\{y\}^\{\\prime\}\}\{\\partial\\theta\_\{k\}\}\\right\)\.\(2\)The second term∂y^′/∂θk\\partial\\hat\{y\}^\{\\prime\}/\\partial\\theta\_\{k\}is the PINN\-specific complication: a second\-order mixed derivative∂2y^/∂θk​∂t\\partial^\{2\}\\hat\{y\}/\\partial\\theta\_\{k\}\\,\\partial tthat requires differentiating the physics\-derivative computation itself with respect to the parameters\. In PyTorch, this term exists only if the physics derivative was computed withcreate\_graph=True; omitting the flag silently drops it, producing wrong gradients without any error message\.

This ODE is the “hello\-world” of physics\-informed learning — small enough to admit closed\-form analysis, yet rich enough to exhibit the two\-derivative coupling above\. A companion paper\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\]develops the analytical derivation of the resulting training gradient for a 1\-3\-3\-1 MLP applied to this problem; the present paper traces the corresponding machine computation inside PyTorch’s AD engine\.

##### From hand\-derived gradients to framework AD\.

Training neural networks to satisfy differential equations predates modern deep learning\. Early work by Dissanayake and Phan\-Thien\[[4](https://arxiv.org/html/2607.13042#bib.bib7)\]and Lagaris et al\.\[[10](https://arxiv.org/html/2607.13042#bib.bib8)\]derived the gradient expressions required for training analytically, for shallow architectures with fixed activation functions; the resulting formulas had to be re\-derived whenever the architecture changed\. The emergence of general\-purpose AD inside deep\-learning frameworks — in particular reverse\-mode AD as exposed by PyTorch\[[11](https://arxiv.org/html/2607.13042#bib.bib3)\]— removed this constraint by computing both the physics derivatives \(such as∂y^/∂t\\partial\\hat\{y\}/\\partial t\) and the parameter gradients∇θL\\nabla\_\{\\theta\}Lgenerically, regardless of network depth or activation choice\. This shift made deep architectures tractable and enabled the modern PINN framework of Raissi et al\.\[[12](https://arxiv.org/html/2607.13042#bib.bib11)\], in which boundary or initial conditions are enforced through additional loss terms, as well as concurrent neural\-network PDE solvers such as the Deep Galerkin method\[[15](https://arxiv.org/html/2607.13042#bib.bib10)\]and the Deep Ritz method\[[5](https://arxiv.org/html/2607.13042#bib.bib9)\]\. Since 2019, PINNs have been applied to fluid dynamics\[[13](https://arxiv.org/html/2607.13042#bib.bib12)\], inverse problems\[[8](https://arxiv.org/html/2607.13042#bib.bib13)\], and a broad range of scientific computing tasks\[[3](https://arxiv.org/html/2607.13042#bib.bib14),[18](https://arxiv.org/html/2607.13042#bib.bib18)\]\.

##### The pedagogical gap\.

The theoretical foundations of AD are well established\. Griewank and Walther\[[7](https://arxiv.org/html/2607.13042#bib.bib1)\]provide the definitive treatment of forward and reverse modes; Baydin et al\.\[[1](https://arxiv.org/html/2607.13042#bib.bib2)\]survey the field for a machine learning audience; standard textbooks\[[6](https://arxiv.org/html/2607.13042#bib.bib5)\]cover backpropagation, the application of reverse\-mode AD to neural network training, following Rumelhart et al\.\[[14](https://arxiv.org/html/2607.13042#bib.bib4)\]\. On the PINN side, introductory treatments such as Katsikis et al\.\[[9](https://arxiv.org/html/2607.13042#bib.bib19)\]explain the loss formulation and training methodology, while peer\-reviewed didactic derivations such as Blechschmidt and Ernst\[[2](https://arxiv.org/html/2607.13042#bib.bib20)\]work through the residual symbolically by chain rule on a single\-hidden\-layer network — stopping short of the AD engine itself\. What is missing is the bridge between the analytical derivation and the machine computation: a single document that traces, at the level of individual operations and explicit numerical values, how the AD engine handles the two\-derivative coupling of[Eqs\.˜1](https://arxiv.org/html/2607.13042#S1.E1)and[2](https://arxiv.org/html/2607.13042#S1.E2)on a multi\-layer network\. Specifically:

- •What data structure does PyTorch build during the forward pass?
- •How does a single backward traversal compute all 22 parameter gradients simultaneously?
- •What doescreate\_graph=Trueactually do, and why does omitting it silently produce wrong gradients?
- •How do the adjoint vectors propagated by PyTorch relate to theP/QP/Qsensitivities derived by hand in\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\]?

Existing references address pieces of this story — AD theory without PINN specifics, PINN formulations without AD internals — but none assembles them into a single, verifiable narrative on a worked example\. Closing this gap prevents silent implementation errors, clarifies the computational cost of physics\-informed training, and provides the foundation for developing techniques that improve convergence and accuracy\.

##### What this paper covers\.

We trace the complete AD pipeline for the same 1\-3\-3\-1 MLP and the same IVPy′\+y=0y^\{\\prime\}\+y=0,y​\(0\)=1y\(0\)=1used in\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\]\. This pairing is small enough that every intermediate quantity — the activations of the forward pass, the 22 parameter gradients of the backward pass, and the extended graph created bycreate\_graph=True— can be displayed explicitly, yet large enough that the two\-derivative coupling of[Eqs\.˜1](https://arxiv.org/html/2607.13042#S1.E1)and[2](https://arxiv.org/html/2607.13042#S1.E2)appears in full generality:

- •[Section˜2](https://arxiv.org/html/2607.13042#S2)introduces tangent and adjoint vectors on a minimal two\-input composition\.
- •[Section˜3](https://arxiv.org/html/2607.13042#S3)traces the forward pass of the 1\-3\-3\-1 MLP and catalogues the computational graph\.
- •[Section˜4](https://arxiv.org/html/2607.13042#S4)traces the backward pass node by node, computing all 22 parameter gradients in one traversal\.
- •[Section˜5](https://arxiv.org/html/2607.13042#S5)explains the graph\-on\-graph: howcreate\_graph=Trueenables correct training gradients, including the product rule andϕ′′\\phi^\{\\prime\\prime\}terms\.
- •[Section˜6](https://arxiv.org/html/2607.13042#S6)addresses common pitfalls, memory, and implementation patterns\.

Every intermediate value is verified against the hand derivations of\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\]and against PyTorch’s output; a companion Jupyter notebook reproduces all calculations\. The physics derivativey^′\\hat\{y\}^\{\\prime\}can also be computed by forward\-mode AD — during the forward pass, a tangent vectora˙\(ℓ\)\\dot\{a\}^\{\(\\ell\)\}is carried alongside each activation and updated by the local Jacobian at each layer, exactly the dual propagation in the forward\-pass tables of\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\]— and PyTorch supports this via[torch\.func\.jvp](https://arxiv.org/html/2607.13042v1/torch.func.jvp)\. We trace the reverse\-mode approach because it is the PyTorch default for PINN training\. This paper does not propose a new PINN methodology, benchmark training strategies, or compare AD systems; its contribution is explanatory\.

## 2Tangents and Adjoints: The Two Faces of the Chain Rule

Automatic differentiation \(AD\) is a family of algorithms that compute derivatives by decomposing a computation into elementary operations and applying the chain rule systematically\. PyTorch records these operations in a data structure called the*computational graph*\. Before tracing derivatives through a neural network, we introduce the forward pass, forward\-mode AD, backward\-mode AD, and the graph\-on\-graph mechanism on a scalar function of two inputs — the smallest example that contains every ingredient PINN training needs\.

### 2\.1A Minimal Example

Lety=f​\(g​\(x1,x2\)\)y=f\(g\(x\_\{1\},x\_\{2\}\)\)with:

g​\(x1,x2\)=2​x1\+x2,f​\(u\)=u2\.g\(x\_\{1\},x\_\{2\}\)=2x\_\{1\}\+x\_\{2\},\\qquad f\(u\)=u^\{2\}\.\(3\)Atx1=2x\_\{1\}=2,x2=1x\_\{2\}=1:

u=g​\(2,1\)=5,y=f​\(5\)=25\.u=g\(2,1\)=5,\\qquad y=f\(5\)=25\.\(4\)The local derivatives are:

∂g∂x1=2,∂g∂x2=1,f′​\(u\)=2​u=10\.\\frac\{\\partial g\}\{\\partial x\_\{1\}\}=2,\\qquad\\frac\{\\partial g\}\{\\partial x\_\{2\}\}=1,\\qquad f^\{\\prime\}\(u\)=2u=10\.\(5\)By the chain rule,∂y/∂x1=10⋅2=20\\partial y/\\partial x\_\{1\}=10\\cdot 2=20and∂y/∂x2=10⋅1=10\\partial y/\\partial x\_\{2\}=10\\cdot 1=10\. We first computey=25y=25and trace what PyTorch records, then compute these derivatives in two ways that correspond to the two modes of automatic differentiation\.

### 2\.2The Forward Pass and Graph 1

The forward pass computesy=25y=25by evaluating the operations in sequence:u=2​x1\+x2=5u=2x\_\{1\}\+x\_\{2\}=5, theny=u2=25y=u^\{2\}=25\. When an input hasrequires\_grad=True, PyTorch builds a computational graph that records this sequence\. Each operation becomes a*node*in the graph — a record containing:

- •the operation performed and the numerical result \(the primal value\),
- •links to the input nodes and the Jacobian entries \(gradient functiongrad\_fn\) required to compute the derivative during the backward pass\.

[Figure˜1](https://arxiv.org/html/2607.13042#S2.F1)shows Graph 1 built during this forward pass\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x1.png)Figure 1:Graph 1 built during the forward pass fory=f​\(g​\(x1,x2\)\)y=f\(g\(x\_\{1\},x\_\{2\}\)\)\. Gray boxes: leaf nodes \(inputs\)\. Blue boxes: operation nodes, each displaying the node name, the operation, the computed primal value, and the associated*gradient function*\. Arrows indicate dependencies, pointing from the output back to the parent nodes that fed the operation\.Graph 1 is the recipe for computing derivatives: it stores the operations and the numerical values needed to evaluate local derivatives during a derivative pass\. No input is perturbed \(as in finite differences\), and no symbolic formula is manipulated — the graph is a trace of elementary operations, each paired with a known, exact derivative rule\.

### 2\.3Forward\-Mode AD: Propagating a Tangent

Forward\-mode AD computes derivatives by propagating a*tangent vector*𝐱˙\\dot\{\\mathbf\{x\}\}from the inputs through the computational graph\. By seeding the input with a vector𝐯=\(x˙1,x˙2\)\\mathbf\{v\}=\(\\dot\{x\}\_\{1\},\\dot\{x\}\_\{2\}\), which can be viewed asd​𝐱/d​λd\\mathbf\{x\}/d\\lambdafor an independent parameterλ\\lambda, the derivative is computed at each node via the rule:

output˙=Jlocal⋅input˙\\dot\{\\text\{output\}\}=J\_\{\\text\{local\}\}\\cdot\\dot\{\\text\{input\}\}\(6\)
This operation is aJacobian–vector product \(JVP\), where them×nm\\times nJacobian left\-multiplies ann×1n\\times 1tangent vector\. This preserves the flow of data, mapping thenninput perturbations tommoutput tangents\. In the transition from inputs to nodeuu, the1×21\\times 2JacobianJu=\[∂u/∂x1,∂u/∂x2\]J\_\{u\}=\[\\partial u/\\partial x\_\{1\},\\partial u/\\partial x\_\{2\}\]multiplies the2×12\\times 1tangent\(x˙1,x˙2\)⊤\(\\dot\{x\}\_\{1\},\\dot\{x\}\_\{2\}\)^\{\\top\}to produce the scalar tangentu˙\\dot\{u\}\.

The final resulty˙\\dot\{y\}represents the directional derivative ofyyin the direction of the seed𝐯\\mathbf\{v\}, scaled by its magnitude\. Specifically, choosing the standard basis vectoreie\_\{i\}as the seed recovers the partial derivative∂y/∂xi\\partial y/\\partial x\_\{i\}\.

[Table˜1](https://arxiv.org/html/2607.13042#S2.T1)traces the two passes needed to obtain both partial derivatives\.

Table 1:Forward\-mode AD requires one pass per input\. Each pass propagates a tangent seed through Graph 1 using the same sequence of JVPs\.[Figure˜2](https://arxiv.org/html/2607.13042#S2.F2)shows the two passes through Graph 1\. Tangent values \(green, bold\) are shown inside each node\. Pass 1 seedsx˙1=1,x˙2=0\\dot\{x\}\_\{1\}=1,\\dot\{x\}\_\{2\}=0and producesy˙=20=∂y/∂x1\\dot\{y\}=20=\\partial y/\\partial x\_\{1\}\. Pass 2 seedsx˙1=0,x˙2=1\\dot\{x\}\_\{1\}=0,\\dot\{x\}\_\{2\}=1and producesy˙=10=∂y/∂x2\\dot\{y\}=10=\\partial y/\\partial x\_\{2\}\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x2.png)

![Refer to caption](https://arxiv.org/html/2607.13042v1/x3.png)

Figure 2:Graph 1 traversed in forward mode\.Top:Pass 1 with seedx˙1=1,x˙2=0\\dot\{x\}\_\{1\}=1,\\dot\{x\}\_\{2\}=0\.Bottom:Pass 2 with seedx˙1=0,x˙2=1\\dot\{x\}\_\{1\}=0,\\dot\{x\}\_\{2\}=1\.
### 2\.4Backward\-Mode AD: Propagating an Adjoint

Backward\-mode AD utilizes the computational graph to compute derivatives by propagating an*adjoint*from the output back to the inputs\. The process begins at the output with a seedy¯=d​y/d​y=1\\bar\{y\}=dy/dy=1\. At each node, the adjoint is propagated backward according to the rule:

input¯=output¯⋅Jlocal\\bar\{\\text\{input\}\}=\\bar\{\\text\{output\}\}\\cdot J\_\{\\text\{local\}\}\(7\)
This is avector–Jacobian product \(VJP\)\. To propagate sensitivities backward, the1×m1\\times madjoint left\-multiplies them×nm\\times nJacobian\. This order is dimensionally required to "pull" the sensitivity of themmoutputs back to thenninputs, resulting in a1×n1\\times nadjoint vector\. Projecting this to our example, the scalar adjointu¯\\bar\{u\}\(a1×11\\times 1vector\) multiplies the1×21\\times 2JacobianJuJ\_\{u\}to simultaneously "distribute" the sensitivity to bothx¯1\\bar\{x\}\_\{1\}andx¯2\\bar\{x\}\_\{2\}in a1×21\\times 2row\.

Here, the overbar notationv¯=∂y/∂v\\bar\{v\}=\\partial y/\\partial vdenotes the*adjoint*ofvv, representing the sensitivity of the final outputyywith respect to the intermediate variablevv\. Unlike forward\-mode, a single backward pass efficiently recovers the gradient of the output with respect to all input parameters simultaneously\.

[Table˜2](https://arxiv.org/html/2607.13042#S2.T2)traces the single pass\.

Table 2:Backward\-mode AD computes both gradients in one pass by propagating an adjoint from output to inputs through Graph 1\.One pass produced both∂y/∂x1=20\\partial y/\\partial x\_\{1\}=20and∂y/∂x2=10\\partial y/\\partial x\_\{2\}=10, plus the intermediate adjointu¯=10\\bar\{u\}=10\.

[Figure˜3](https://arxiv.org/html/2607.13042#S2.F3)shows the same Graph 1 traversed in backward mode\. The adjoint seedy¯=1\\bar\{y\}=1starts at the output node and flows right to left\. At each node, the incoming adjoint is multiplied by the local derivative on each outgoing edge to produce the adjoint at the next node:u¯=y¯⋅10=10\\bar\{u\}=\\bar\{y\}\\cdot 10=10, thenx¯1=u¯⋅2=20\\bar\{x\}\_\{1\}=\\bar\{u\}\\cdot 2=20andx¯2=u¯⋅1=10\\bar\{x\}\_\{2\}=\\bar\{u\}\\cdot 1=10\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x4.png)Figure 3:Graph 1 traversed in backward mode\. Adjoint values \(red, bold\) are shown inside each node\.##### The key distinction\.

Both modes apply the same local derivatives \(f′=10f^\{\\prime\}=10,∂g/∂x1=2\\partial g/\\partial x\_\{1\}=2,∂g/∂x2=1\\partial g/\\partial x\_\{2\}=1\)\. They differ in direction and cost:

For any gradient\-based optimization involving a scalar loss and a high\-dimensional parameter space \(m=1,n≫1m=1,n\\gg 1\), backward\-mode AD is significantly more efficient as it computes the full gradient∇y∈ℝn\\nabla y\\in\\mathbb\{R\}^\{n\}in a single backward pass\. This computational advantage is the reason PyTorch adopts backward\-mode as its default autodiff engine and provides theloss\.backward\(\)interface for gradient calculation\.

### 2\.5Second Order Derivative and the Graph\-on\-Graph Mechanism

The backward pass computedx¯1=20\\bar\{x\}\_\{1\}=20through VJP operations:u¯=y¯⋅2​u=10\\bar\{u\}=\\bar\{y\}\\cdot 2u=10, thenx¯1=u¯⋅2=20\\bar\{x\}\_\{1\}=\\bar\{u\}\\cdot 2=20\. These VJP operations are themselves arithmetic\. Withcreate\_graph=True, PyTorch records them in a second graph \(Graph 2\), making it possible to differentiatex¯1\\bar\{x\}\_\{1\}with respect to the original inputs\.[Figure˜4](https://arxiv.org/html/2607.13042#S2.F4)shows both graphs: Graph 1 \(top\) built during the forward pass, and Graph 2 \(bottom\) built during the first backward pass\. The dashed cross\-link records that the VJP atyyreadsu=5u=5from Graph 1\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x5.png)Figure 4:The two graphs built during forward and backward passes\.Top:Graph 1, built during the forward pass whenrequires\_grad=True; gray boxes are leaf nodes, blue boxes are operation nodes\.Bottom:Graph 2, built during the first backward pass only whencreate\_graph=Trueis set\.Withoutcreate\_graph=True, the VJP operations execute but are not recorded\. The mixed derivative becomes∂x¯1/∂x2=∂\(20\)/∂x2=0\\partial\\bar\{x\}\_\{1\}/\\partial x\_\{2\}=\\partial\(20\)/\\partial x\_\{2\}=0\(wrong; correct answer is 4\)\.

[Figure˜5](https://arxiv.org/html/2607.13042#S2.F5)shows the traversal propagating the adjoint fromx¯1\\bar\{x\}\_\{1\}node with seed11back tox2x\_\{2\}as a sequence of node visits, with the adjoint accumulating from 1 to 4 as it flows through the combined graphG​r​a​p​h​1\+G​r​a​p​h​2Graph\\,1\+Graph\\,2\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x6.png)Figure 5:Traversal path for computing∂x¯1/∂x2=∂2y∂x2​∂x1\\partial\\bar\{x\}\_\{1\}/\\partial x\_\{2\}=\\dfrac\{\\partial^\{2\}y\}\{\{\\partial x\_\{2\}\}\{\\partial x\_\{1\}\}\}\.##### The PINN analogy\.

The roles in PINN training for the initial value problem presented in[Section˜3\.1](https://arxiv.org/html/2607.13042#S3.SS1)map directly onto this example:

### 2\.6Verification in PyTorch

To verify the theoretical framework, the following listings demonstrate PyTorch’s tiered approach to automatic differentiation, progressing from a standard forward pass to the computation of second\-order mixed derivatives\. While Level 0 ignores gradient tracking, Level 1 allows for first\-order gradients for the leaf tensors havingrequires\_grad=True\. The use ofretain\_graph=Truein the first backward AD traversal is essential because PyTorch’s default behavior is to free the intermediate buffers once a backward pass is completed; without this flag, the computation graph would be destroyed after calculatingx¯1\\bar\{x\}\_\{1\}, making the subsequent calculation ofx¯2\\bar\{x\}\_\{2\}impossible\. Level 2 utilizescreate\_graph=Trueto build a graph capable of tracking the second order derivatives — the derivatives ofx¯1=∂y/∂x1\\bar\{x\}\_\{1\}=\\partial y/\\partial x\_\{1\}\.

Listing 1:Level 0: norequires\_grad, no derivatives\.1importtorch

2

3x1=torch\.tensor\(2\.0\)

4x2=torch\.tensor\(1\.0\)

5

6u=2\*x1\+x2

7y=u\*\*2

8

9print\(y\.item\(\)\)

10print\(y\.grad\_fn\)

Listing 2:Level 1: first derivatives work; mixed derivative fails\.1x1=torch\.tensor\(2\.0,requires\_grad=True\)

2x2=torch\.tensor\(1\.0,requires\_grad=True\)

3

4u=2\*x1\+x2

5y=u\*\*2

6

7

8bar\_x1=torch\.autograd\.grad\(y,x1,

9retain\_graph=True\)\[0\]

10bar\_x2=torch\.autograd\.grad\(y,x2\)\[0\]

11

12print\(bar\_x1\.item\(\)\)

13print\(bar\_x2\.item\(\)\)

14print\(bar\_x1\.grad\_fn\)

15

16

17try:

18torch\.autograd\.grad\(bar\_x1,x2\)

19exceptRuntimeErrorase:

20print\(e\)

21

22

Listing 3:Level 2:create\_graph=Trueenables the mixed derivative\.1x1=torch\.tensor\(2\.0,requires\_grad=True\)

2x2=torch\.tensor\(1\.0,requires\_grad=True\)

3

4u=2\*x1\+x2

5y=u\*\*2

6

7

8bar\_x1=torch\.autograd\.grad\(y,x1,

9create\_graph=True\)\[0\]

10

11print\(bar\_x1\.item\(\)\)

12print\(type\(bar\_x1\.grad\_fn\)\.\_\_name\_\_\)

13

14mixed=torch\.autograd\.grad\(bar\_x1,x2\)\[0\]

15print\(mixed\.item\(\)\)

With the mechanics of higher\-order AD established, we are ready to apply these concepts to the development of a Physics\-Informed Neural Network \(PINN\) model designed to solve a fundamental initial value problem\.

## 3The Forward Pass

We now move from the two\-input composition to a real neural network\. This section defines the problem and the network that will serve as our running example for the rest of the paper, computes the forward pass with explicit numerical values, and traces the computational graph that PyTorch builds along the way\.

### 3\.1The Problem

We seek to approximate the solution of the initial value problem

y′​\(t\)\+y​\(t\)=0,y​\(0\)=1,y^\{\\prime\}\(t\)\+y\(t\)=0,\\qquad y\(0\)=1,\(8\)whose analytical solution isy​\(t\)=e−ty\(t\)=e^\{\-t\}\. A Physics\-Informed Neural Network \(PINN\) approximatesy​\(t\)y\(t\)by a neural networky^​\(t;θ\)\\hat\{y\}\(t;\\theta\)and trains it by minimizing a loss function that penalizes the ODE residualR​\(t\)=y^′​\(t\)\+y^​\(t\)R\(t\)=\\hat\{y\}^\{\\prime\}\(t\)\+\\hat\{y\}\(t\)at collocation points and the initial condition errory^​\(0\)−1\\hat\{y\}\(0\)\-1\. The details of the loss will be needed in[Section˜5](https://arxiv.org/html/2607.13042#S5); for now, the key point is that the loss depends on bothy^\\hat\{y\}and its derivativey^′=d​y^/d​t\\hat\{y\}^\{\\prime\}=d\\hat\{y\}/dt, and the AD engine must handle both\.

### 3\.2The Network

The network is a1\-3\-3\-1 multilayer perceptron: one input \(tt\), two hidden layers with 3 neurons each usingtanh\\tanhactivation, and one linear output neuron producingy^​\(t\)\\hat\{y\}\(t\)\. The forward pass through layerℓ\\ellis:

z\(ℓ\)=W\(ℓ\)​a\(ℓ−1\)\+b\(ℓ\),a\(ℓ\)=ϕ​\(z\(ℓ\)\),z^\{\(\\ell\)\}=W^\{\(\\ell\)\}a^\{\(\\ell\-1\)\}\+b^\{\(\\ell\)\},\\qquad a^\{\(\\ell\)\}=\\phi\(z^\{\(\\ell\)\}\),\(9\)whereϕ=tanh\\phi=\\tanhfor hidden layers andϕ=identity\\phi=\\text\{identity\}for the output layer\. The network has 22 trainable parameters:θ=\{W\(ℓ\),b\(ℓ\)\}ℓ=13\\theta=\\\{W^\{\(\\ell\)\},b^\{\(\\ell\)\}\\\}\_\{\\ell=1\}^\{3\}\.

We fix all parameters to enable hand verification\.[Table˜3](https://arxiv.org/html/2607.13042#S3.T3)lists the complete parameter set\.

Table 3:Fixed parameters for the 1\-3\-3\-1 MLP \(22 parameters total\)\.
### 3\.3Forward Pass att=0\.5t=0\.5

We trace the forward pass at the collocation pointtc=0\.5t\_\{c\}=0\.5\. The input isa\(0\)=0\.5a^\{\(0\)\}=0\.5\.

##### Layer 1\.

z\(1\)=W\(1\)​\(0\.5\)\+b\(1\)=\(0\.1−0\.1−0\.25\+0\.30\.4−0\.2\)=\(0\.00\.050\.2\),a\(1\)=tanh⁡\(z\(1\)\)=\(0\.00000\.05000\.1974\)\.z^\{\(1\)\}=W^\{\(1\)\}\(0\.5\)\+b^\{\(1\)\}=\\begin\{pmatrix\}0\.1\-0\.1\\\\ \-0\.25\+0\.3\\\\ 0\.4\-0\.2\\end\{pmatrix\}=\\begin\{pmatrix\}0\.0\\\\ 0\.05\\\\ 0\.2\\end\{pmatrix\},\\qquad a^\{\(1\)\}=\\tanh\(z^\{\(1\)\}\)=\\begin\{pmatrix\}0\.0000\\\\ 0\.0500\\\\ 0\.1974\\end\{pmatrix\}\.\(10\)

##### Layer 2\.

z\(2\)=W\(2\)​a\(1\)\+b\(2\)=\(0\.2837−0\.16900\.4547\),a\(2\)=tanh⁡\(z\(2\)\)=\(0\.2763−0\.16730\.4257\)\.z^\{\(2\)\}=W^\{\(2\)\}a^\{\(1\)\}\+b^\{\(2\)\}=\\begin\{pmatrix\}0\.2837\\\\ \-0\.1690\\\\ 0\.4547\\end\{pmatrix\},\\qquad a^\{\(2\)\}=\\tanh\(z^\{\(2\)\}\)=\\begin\{pmatrix\}0\.2763\\\\ \-0\.1673\\\\ 0\.4257\\end\{pmatrix\}\.\(11\)To verify one entry:z1\(2\)=0\.1​\(0\.0000\)\+\(−0\.3\)​\(0\.0500\)\+0\.5​\(0\.1974\)\+0\.2=0−0\.0150\+0\.0987\+0\.2=0\.2837z^\{\(2\)\}\_\{1\}=0\.1\(0\.0000\)\+\(\-0\.3\)\(0\.0500\)\+0\.5\(0\.1974\)\+0\.2=0\-0\.0150\+0\.0987\+0\.2=0\.2837\.

##### Layer 3 \(output\)\.

y^=W\(3\)​a\(2\)\+b\(3\)=0\.9​\(0\.2763\)\+\(−0\.6\)​\(−0\.1673\)\+0\.3​\(0\.4257\)−0\.3=0\.1768\.\\hat\{y\}=W^\{\(3\)\}a^\{\(2\)\}\+b^\{\(3\)\}=0\.9\(0\.2763\)\+\(\-0\.6\)\(\-0\.1673\)\+0\.3\(0\.4257\)\-0\.3=0\.1768\.\(12\)
[Table˜4](https://arxiv.org/html/2607.13042#S3.T4)collects the complete forward pass\.

Table 4:Forward pass fort=0\.5t=0\.5\.

### 3\.4What PyTorch Records

Two points deserve emphasis\. First, the graph records*elementary operations*, not neural network layers\. A single[nn\.Linear](https://arxiv.org/html/2607.13042v1/nn.Linear)layer produces at least two operation nodes: one for the matrix multiply, one for the bias addition\. Second, the leaf nodes now include both the inputttand the parameter tensorsW\(ℓ\)W^\{\(\\ell\)\},b\(ℓ\)b^\{\(\\ell\)\}\. In PyTorch, parameters are leaves with[requires\_grad=True](https://arxiv.org/html/2607.13042v1/requires_grad=True)by default; the inputttbecomes a tracked leaf only if the user sets[requires\_grad=True](https://arxiv.org/html/2607.13042v1/requires_grad=True)explicitly — which is necessary for computingy^′\\hat\{y\}^\{\\prime\}\.

[Table˜5](https://arxiv.org/html/2607.13042#S3.T5)lists every operation node that PyTorch creates during the forward pass of[Table˜4](https://arxiv.org/html/2607.13042#S3.T4)\. For each node, the table shows the tensor it produces, the[grad\_fn](https://arxiv.org/html/2607.13042v1/grad_fn)that PyTorch attaches to it, and the tensors it saves for the backward pass\.

Table 5:Forward pass trace \(Graph 1 construction\) fort=0\.5↦y^=0\.1768t=0\.5\\mapsto\\hat\{y\}=0\.1768\. Each primal node materializes an intermediate value and buffers the local Jacobian factors required for the first\-order backward pass\.##### Three observations\.

1. 1\.The saved tensors are the chain rule factors\.Consider the nodeL3 matmul, which computesz\(3\)=a\(2\)​W\(3\)z^\{\(3\)\}=a^\{\(2\)\}W^\{\(3\)\}: it saves the input activationa\(2\)a^\{\(2\)\}and the weight matrixW\(3\)W^\{\(3\)\}\. In the backward pass, these tensors serve as the local Jacobian factors for two distinct VJPs\. The incoming adjointz¯\(3\)\\bar\{z\}^\{\(3\)\}is multiplied bya\(2\)a^\{\(2\)\}to produce the weight gradientW¯\(3\)=∂L/∂W\(3\)\\bar\{W\}^\{\(3\)\}=\\partial L/\\partial W^\{\(3\)\}for the optimizer, and multiplied byW\(3\)W^\{\(3\)\}to propagate the sensitivity back as the activation adjoint of the previous layera¯\(2\)=∂L/∂a\(2\)\\bar\{a\}^\{\(2\)\}=\\partial L/\\partial a^\{\(2\)\}\.
2. 2\.Tanh saves its output, not its input\.The VJP oftanh\\tanhrequires the derivativeϕ′​\(z\)=1−tanh2⁡\(z\)=1−a2\\phi^\{\\prime\}\(z\)=1\-\\tanh^\{2\}\(z\)=1\-a^\{2\}, which can be reconstructed from the outputaaalone\. For example, at nodeL2 tanh:ϕ′​\(z\(2\)\)=1−\(a\(2\)\)2=\(0\.9237,0\.9720,0\.8188\)\\phi^\{\\prime\}\(z^\{\(2\)\}\)=1\-\(a^\{\(2\)\}\)^\{2\}=\(0\.9237,\\;0\.9720,\\;0\.8188\)\. This is a deliberate design choice: savingaainstead ofzzavoids recomputing the expensive exponential functions withintanh\\tanhduring the backward pass\.
3. 3\.Addition nodes record nothing\.Nodes likeL3 addrecord nothing\(none\)\. Since the partial derivative of an additionz=u\+bz=u\+bis simply11with respect to both inputs \(whereuurepresents the product term andbbthe bias\), the incoming adjointz¯\\bar\{z\}is passed back unchanged:u¯=z¯⋅1\\bar\{u\}=\\bar\{z\}\\cdot 1to continue the backward flow, andb¯=z¯⋅1\\bar\{b\}=\\bar\{z\}\\cdot 1to serve as the gradient for the optimizer\. No records from the forward pass are needed to compute these adjoints and no actual calculation is required, which minimizes both memory overhead and computational effort\.

### 3\.5The Complete DAG

[Figure˜6](https://arxiv.org/html/2607.13042#S3.F6)assembles all the nodes from[Table˜5](https://arxiv.org/html/2607.13042#S3.T5)into the full directed acyclic graph\. The gray boxes on the left are the leaf nodes \(input and parameters\); the white boxes are the operation nodes\. Arrows indicate dependencies — the direction the adjoint will flow when we traverse this graph backward in[Section˜4](https://arxiv.org/html/2607.13042#S4)\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x7.png)Figure 6:Computational graph fort=0\.5↦y^=0\.1768t=0\.5\\mapsto\\hat\{y\}=0\.1768\. Gray boxes: leaf nodes\. White boxes: operation nodes\. Arrows point in the direction of adjoint flow\.
### 3\.6Inspecting the Graph in PyTorch

The computational graph is not just a conceptual diagram — it is a data structure that PyTorch builds in memory and that the user can inspect directly\. The following listings provide a step\-by\-step walkthrough of the graph’s instantiation and the programmatic inspection of its internal operation nodes\.

##### Defining the network\.

We construct the neural network model and set the fixed parameters from[Table˜3](https://arxiv.org/html/2607.13042#S3.T3):

Listing 4:Network definition and parameter initialization\.1importtorch

2importtorch\.nnasnn

3

4net=nn\.Sequential\(

5nn\.Linear\(1,3\),nn\.Tanh\(\),

6nn\.Linear\(3,3\),nn\.Tanh\(\),

7nn\.Linear\(3,1\),

8\)

9

10withtorch\.no\_grad\(\):

11net\[0\]\.weight\.copy\_\(torch\.tensor\(

12\[\[0\.2\],\[\-0\.5\],\[0\.8\]\]\)\)

13net\[0\]\.bias\.copy\_\(torch\.tensor\(

14\[\-0\.1,0\.3,\-0\.2\]\)\)

15net\[2\]\.weight\.copy\_\(torch\.tensor\(

16\[\[0\.1,\-0\.3,0\.5\],

17\[0\.6,0\.2,\-0\.4\],

18\[\-0\.2,0\.7,0\.1\]\]\)\)

19net\[2\]\.bias\.copy\_\(torch\.tensor\(

20\[0\.2,\-0\.1,0\.4\]\)\)

21net\[4\]\.weight\.copy\_\(torch\.tensor\(

22\[\[0\.9,\-0\.6,0\.3\]\]\)\)

23net\[4\]\.bias\.copy\_\(torch\.tensor\(\[\-0\.3\]\)\)

##### Forward pass and graph inspection\.

Running the forward pass builds the graph\. We then inspect the output node and its parents:

Listing 5:Forward pass and node inspection\.1t=torch\.tensor\(\[\[0\.5\]\],requires\_grad=True\)

2y\_hat=net\(t\)

3

4print\(y\_hat\.item\(\)\)

5

6

7print\(type\(y\_hat\.grad\_fn\)\.\_\_name\_\_\)

8

9

10print\(\[\(type\(p\[0\]\)\.\_\_name\_\_,p\[1\]\)

11forpiny\_hat\.grad\_fn\.next\_functions\]\)

12

13

14print\(t\.grad\_fn,t\.requires\_grad\)

15

The output nodeAddBackward0represents the final addition in the network\. Itsnext\_functionsattribute contains a list of its parent nodes: theL3 matmuland the bias leaf\. These are stored as tuples\(function, index\), where the index \(0\) specifies which output of the parent node feeds the current operation—a detail required for operations that return multiple tensors, which is not the case here\. The loop iterates through this list to propagate gradients back through all input branches\. While the bias appears asAccumulateGradto manage its update, the inputttreturnsNonefor itsgrad\_fnbecause it is a user\-defined root where the graph terminates\.

##### Walking the full graph\.

Following the first parent at each node traces the path from the output back to the input:

Listing 6:Walking the computational graph from output to input\.1node=y\_hat\.grad\_fn

2whilenodeisnotNone:

3print\(type\(node\)\.\_\_name\_\_\)

4parents=node\.next\_functions

5node=parents\[0\]\[0\]ifparentselseNone

6

7

8

9

10

11

12

13

14

The sequence matches[Table˜5](https://arxiv.org/html/2607.13042#S3.T5)exactly: eight operation nodes, terminating at the leaf\. This is the recipe thatloss\.backward\(\)will follow in[Section˜4](https://arxiv.org/html/2607.13042#S4)— the graph stores the sequence of VJP operations, not the gradients themselves\.

## 4Reverse\-Mode AD

We now traverse the graph from[Figure˜6](https://arxiv.org/html/2607.13042#S3.F6)backward, propagating adjoint vectors from the output to the leaves\. By the end of this single traversal, we will have computed∂y^/∂θk\\partial\\hat\{y\}/\\partial\\theta\_\{k\}for all 22 parameters simultaneously\.

### 4\.1Starting Point: The Adjoint Seed

To illustrate reverse\-mode in isolation, we first differentiatey^\\hat\{y\}itself \(not the loss\)\. The adjoint seed is:

y^¯=∂y^∂y^=1\.\\bar\{\\hat\{y\}\}=\\frac\{\\partial\\hat\{y\}\}\{\\partial\\hat\{y\}\}=1\.\(13\)This single number enters the graph at the output node and flows backward\. At each operation node, the VJP rule uses the tensors saved during the forward pass \([Table˜5](https://arxiv.org/html/2607.13042#S3.T5)\) to propagate the adjoint one step further\.

### 4\.2Constructing VJP Nodes During the Backward Pass

Following the conceptual framework of Graph 2 established in[Section˜2\.5](https://arxiv.org/html/2607.13042#S2.SS5), we now trace the backward pass node\-by\-node, formalizing the record of each Vector\-Jacobian Product \(VJP\) as an explicit node\.

#### 4\.2\.1VJP L3 Add

In the primal forward pass, the output layer computesy^=m\(3\)\+b\(3\)\\hat\{y\}=m^\{\(3\)\}\+b^\{\(3\)\}\. The nodeVJP L3 Addreceives the incoming adjointy^¯=1∈ℝ\\bar\{\\hat\{y\}\}=1\\in\\mathbb\{R\}, which serves as the seed for the backward trace\. Since the output activation is the identity,z¯\(3\)=y^¯\\bar\{z\}^\{\(3\)\}=\\bar\{\\hat\{y\}\}\.

The edges connecting the output node to its primal parents \(m\(3\)m^\{\(3\)\}andb\(3\)b^\{\(3\)\}\) represent the local Jacobians∂y^/∂m\(3\)=1\\partial\\hat\{y\}/\\partial m^\{\(3\)\}=1and∂y^/∂b\(3\)=1\\partial\\hat\{y\}/\\partial b^\{\(3\)\}=1\. Applying the VJP operation—multiplying the incoming adjoint by these local gradients—yields:

m¯\(3\)=y^¯⋅1=1,b¯\(3\)=y^¯⋅1=1\.\\bar\{m\}^\{\(3\)\}=\\bar\{\\hat\{y\}\}\\cdot 1=1,\\qquad\\bar\{b\}^\{\(3\)\}=\\bar\{\\hat\{y\}\}\\cdot 1=1\.\(14\)
The resulting adjointb¯\(3\)∈ℝ\\bar\{b\}^\{\(3\)\}\\in\\mathbb\{R\}is the terminal gradient for the output bias, whilem¯\(3\)∈ℝ\\bar\{m\}^\{\(3\)\}\\in\\mathbb\{R\}is propagated backward to serve as the incoming adjoint for theVJP L3 Mmnode\.

#### 4\.2\.2VJP L3 Mm

In the primal pass, the linear transformation ism\(3\)=a\(2\)​\(W\(3\)\)Tm^\{\(3\)\}=a^\{\(2\)\}\(W^\{\(3\)\}\)^\{T\}\. The nodeVJP L3 Mmreceives the incoming adjointm¯\(3\)=1\\bar\{m\}^\{\(3\)\}=1\. To propagate the adjoint, the node evaluates the VJPs with respect to its primal parents,a\(2\)a^\{\(2\)\}andW\(3\)W^\{\(3\)\}:

a¯\(2\)\\displaystyle\\bar\{a\}^\{\(2\)\}=m¯\(3\)​W\(3\)\\displaystyle=\\bar\{m\}^\{\(3\)\}W^\{\(3\)\}\(15\)=\(1\)​\(0\.9−0\.60\.3\)=\(0\.9−0\.60\.3\)∈ℝ1×3,\\displaystyle=\(1\)\\begin\{pmatrix\}0\.9&\-0\.6&0\.3\\end\{pmatrix\}=\\begin\{pmatrix\}0\.9&\-0\.6&0\.3\\end\{pmatrix\}\\in\\mathbb\{R\}^\{1\\times 3\},W¯\(3\)\\displaystyle\\bar\{W\}^\{\(3\)\}=\(m¯\(3\)\)T​a\(2\)\\displaystyle=\(\\bar\{m\}^\{\(3\)\}\)^\{T\}a^\{\(2\)\}\(16\)=\(1\)T​\(0\.2763−0\.16730\.4257\)=\(0\.2763−0\.16730\.4257\)∈ℝ1×3\.\\displaystyle=\(1\)^\{T\}\\begin\{pmatrix\}0\.2763&\-0\.1673&0\.4257\\end\{pmatrix\}=\\begin\{pmatrix\}0\.2763&\-0\.1673&0\.4257\\end\{pmatrix\}\\in\\mathbb\{R\}^\{1\\times 3\}\.
The row vectora¯\(2\)\\bar\{a\}^\{\(2\)\}represents the sensitivity∂y^/∂a\(2\)\\partial\\hat\{y\}/\\partial a^\{\(2\)\}—the gradient of the output with respect to each of the three hidden neurons in the second layer\. This vector is propagated backward to serve as the incoming adjoint for theVJP L2 Tanhnode\. Simultaneously,W¯\(3\)\\bar\{W\}^\{\(3\)\}provides the terminal gradients for the layer weights, which are stored for the optimizer\.

#### 4\.2\.3VJP L2 Tanh

In the primal pass, the activation layer computesa\(2\)=tanh⁡\(z\(2\)\)a^\{\(2\)\}=\\tanh\(z^\{\(2\)\}\)\. The nodeVJP L2 Tanhreceives the incoming adjointa¯\(2\)∈ℝ1×3\\bar\{a\}^\{\(2\)\}\\in\\mathbb\{R\}^\{1\\times 3\}\. Because thetanh\\tanhfunction is applied element\-wise, the Jacobian is a diagonal matrix of local derivativesϕ′​\(z\(2\)\)=1−\(a\(2\)\)2\\phi^\{\\prime\}\(z^\{\(2\)\}\)=1\-\(a^\{\(2\)\}\)^\{2\}\. In Autograd, this VJP is efficiently computed using the Hadamard product:

z¯\(2\)\\displaystyle\\bar\{z\}^\{\(2\)\}=a¯\(2\)⊙\(1−\(a\(2\)\)2\)\\displaystyle=\\bar\{a\}^\{\(2\)\}\\odot\\bigl\(1\-\(a^\{\(2\)\}\)^\{2\}\\bigr\)\(17\)=\(0\.9−0\.60\.3\)⊙\(0\.92370\.97200\.8188\)\\displaystyle=\\begin\{pmatrix\}0\.9&\-0\.6&0\.3\\end\{pmatrix\}\\odot\\begin\{pmatrix\}0\.9237&0\.9720&0\.8188\\end\{pmatrix\}=\(0\.8314−0\.58320\.2456\)∈ℝ1×3\.\\displaystyle=\\begin\{pmatrix\}0\.8314&\-0\.5832&0\.2456\\end\{pmatrix\}\\in\\mathbb\{R\}^\{1\\times 3\}\.
The resulting row vectorz¯\(2\)\\bar\{z\}^\{\(2\)\}represents the sensitivity of the outputy^\\hat\{y\}with respect to the pre\-activations\. This adjoint is propagated backward to theVJP L2 Addnode\. Furthermore, because the biasb\(2\)b^\{\(2\)\}enters the primal pass additively \(z\(2\)=m\(2\)\+b\(2\)z^\{\(2\)\}=m^\{\(2\)\}\+b^\{\(2\)\}\), the local Jacobian∂z\(2\)/∂b\(2\)\\partial z^\{\(2\)\}/\\partial b^\{\(2\)\}is the identity\. Consequently, these entries also provide the terminal gradients for the layer biases:b¯\(2\)=z¯\(2\)\\bar\{b\}^\{\(2\)\}=\\bar\{z\}^\{\(2\)\}\.

#### 4\.2\.4VJP L2 Add

In the primal pass, the pre\-activation is formed byz\(2\)=m\(2\)\+b\(2\)z^\{\(2\)\}=m^\{\(2\)\}\+b^\{\(2\)\}\. The nodeVJP L2 Addreceives the incoming adjointz¯\(2\)∈ℝ1×3\\bar\{z\}^\{\(2\)\}\\in\\mathbb\{R\}^\{1\\times 3\}\. Because the Jacobian of a sum is the identity matrix, the addition node acts as a distributor, passing the adjoint signal unchanged to both the matrix\-product termm\(2\)m^\{\(2\)\}and the biasb\(2\)b^\{\(2\)\}:

m¯\(2\)=z¯\(2\),b¯\(2\)=\(z¯\(2\)\)⊤=\(0\.8314−0\.58320\.2456\)∈ℝ3×1\.\\bar\{m\}^\{\(2\)\}=\\bar\{z\}^\{\(2\)\},\\qquad\\bar\{b\}^\{\(2\)\}=\(\\bar\{z\}^\{\(2\)\}\)^\{\\top\}=\\begin\{pmatrix\}0\.8314\\\\ \-0\.5832\\\\ 0\.2456\\end\{pmatrix\}\\in\\mathbb\{R\}^\{3\\times 1\}\.\(18\)
We transpose the result forb¯\(2\)\\bar\{b\}^\{\(2\)\}to ensure the adjoint matches the column\-vector orientation of the bias parameters as defined in[Table˜3](https://arxiv.org/html/2607.13042#S3.T3)\. The row vectorm¯\(2\)\\bar\{m\}^\{\(2\)\}is then propagated backward to theVJP L2 Mmnode to continue the chain\.

#### 4\.2\.5VJP L2 Mm

In the primal pass, the hidden layer transformation ism\(2\)=a\(1\)​\(W\(2\)\)⊤m^\{\(2\)\}=a^\{\(1\)\}\(W^\{\(2\)\}\)^\{\\top\}\. The nodeVJP L2 Mmreceives the incoming sensitivitym¯\(2\)∈ℝ1×3\\bar\{m\}^\{\(2\)\}\\in\\mathbb\{R\}^\{1\\times 3\}from the addition node\. To propagate the adjoint signal, the node evaluates the VJPs with respect to the activations and weights:

a¯\(1\)\\displaystyle\\bar\{a\}^\{\(1\)\}=m¯\(2\)​W\(2\)\\displaystyle=\\bar\{m\}^\{\(2\)\}W^\{\(2\)\}\(19\)=\(0\.8314−0\.58320\.2456\)​\(0\.1−0\.30\.50\.60\.2−0\.4−0\.20\.70\.1\)\\displaystyle=\\begin\{pmatrix\}0\.8314&\-0\.5832&0\.2456\\end\{pmatrix\}\\begin\{pmatrix\}0\.1&\-0\.3&0\.5\\\\ 0\.6&0\.2&\-0\.4\\\\ \-0\.2&0\.7&0\.1\\end\{pmatrix\}=\(−0\.3159−0\.19390\.6734\)∈ℝ1×3,\\displaystyle=\\begin\{pmatrix\}\-0\.3159&\-0\.1939&0\.6734\\end\{pmatrix\}\\in\\mathbb\{R\}^\{1\\times 3\},W¯\(2\)\\displaystyle\\bar\{W\}^\{\(2\)\}=\(m¯\(2\)\)⊤​a\(1\)\\displaystyle=\(\\bar\{m\}^\{\(2\)\}\)^\{\\top\}a^\{\(1\)\}\(20\)=\(0\.8314−0\.58320\.2456\)​\(0\.00000\.05000\.1974\)\\displaystyle=\\begin\{pmatrix\}0\.8314\\\\ \-0\.5832\\\\ 0\.2456\\end\{pmatrix\}\\begin\{pmatrix\}0\.0000&0\.0500&0\.1974\\end\{pmatrix\}=\(0\.00000\.04160\.16410\.0000−0\.0292−0\.11510\.00000\.01230\.0485\)∈ℝ3×3\.\\displaystyle=\\begin\{pmatrix\}0\.0000&0\.0416&0\.1641\\\\ 0\.0000&\-0\.0292&\-0\.1151\\\\ 0\.0000&0\.0123&0\.0485\\end\{pmatrix\}\\in\\mathbb\{R\}^\{3\\times 3\}\.
The row vectora¯\(1\)\\bar\{a\}^\{\(1\)\}is passed backward as the incoming sensitivity for theVJP L1 Tanhnode\. The matrixW¯\(2\)\\bar\{W\}^\{\(2\)\}provides the terminal gradients for the second\-layer weights\. By constructingW¯\(2\)\\bar\{W\}^\{\(2\)\}as the outer product of the transposed sensitivity and the activation row, the AD engine ensures the gradient dimensions are consistent with the3×33\\times 3parameter matrixW\(2\)W^\{\(2\)\}defined in[Table˜3](https://arxiv.org/html/2607.13042#S3.T3)\.

#### 4\.2\.6VJP L1 Tanh

In the primal pass, the first activation layer computesa\(1\)=tanh⁡\(z\(1\)\)a^\{\(1\)\}=\\tanh\(z^\{\(1\)\}\)\. The nodeVJP L1 Tanhreceives the incoming sensitivitya¯\(1\)∈ℝ1×3\\bar\{a\}^\{\(1\)\}\\in\\mathbb\{R\}^\{1\\times 3\}\. As with the second layer, the element\-wise nature of the activation function results in a diagonal Jacobian, and the VJP is evaluated via the Hadamard product:

z¯\(1\)\\displaystyle\\bar\{z\}^\{\(1\)\}=a¯\(1\)⊙\(1−\(a\(1\)\)2\)\\displaystyle=\\bar\{a\}^\{\(1\)\}\\odot\\bigl\(1\-\(a^\{\(1\)\}\)^\{2\}\\bigr\)\(21\)=\(−0\.3159−0\.19390\.6734\)⊙\(1\.00000\.99750\.9610\)\\displaystyle=\\begin\{pmatrix\}\-0\.3159&\-0\.1939&0\.6734\\end\{pmatrix\}\\odot\\begin\{pmatrix\}1\.0000&0\.9975&0\.9610\\end\{pmatrix\}=\(−0\.3159−0\.19340\.6471\)∈ℝ1×3\.\\displaystyle=\\begin\{pmatrix\}\-0\.3159&\-0\.1934&0\.6471\\end\{pmatrix\}\\in\\mathbb\{R\}^\{1\\times 3\}\.
The resulting row vectorz¯\(1\)\\bar\{z\}^\{\(1\)\}is propagated to theVJP L1 Addnode\. These entries represent the sensitivity of the outputy^\\hat\{y\}with respect to the first\-layer pre\-activations\. Because the biasb\(1\)b^\{\(1\)\}enters the primal pass additively, these sensitivities also constitute the terminal gradients for the first\-layer biases, whereb¯\(1\)=\(z¯\(1\)\)⊤∈ℝ3×1\\bar\{b\}^\{\(1\)\}=\(\\bar\{z\}^\{\(1\)\}\)^\{\\top\}\\in\\mathbb\{R\}^\{3\\times 1\}\.

#### 4\.2\.7VJP L1 Add

In the primal pass, the first\-layer pre\-activation isz\(1\)=m\(1\)\+b\(1\)z^\{\(1\)\}=m^\{\(1\)\}\+b^\{\(1\)\}\. The nodeVJP L1 Addreceives the incoming sensitivityz¯\(1\)\\bar\{z\}^\{\(1\)\}\. As addition act as a distributor in the adjoint graph, the sensitivity flows unchanged to both the matrix\-product termm\(1\)m^\{\(1\)\}and the biasb\(1\)b^\{\(1\)\}:

m¯\(1\)=z¯\(1\)=\(−0\.3159−0\.19340\.6471\)∈ℝ1×3\.\\bar\{m\}^\{\(1\)\}=\\bar\{z\}^\{\(1\)\}=\\begin\{pmatrix\}\-0\.3159&\-0\.1934&0\.6471\\end\{pmatrix\}\\in\\mathbb\{R\}^\{1\\times 3\}\.\(22\)
The row vectorm¯\(1\)\\bar\{m\}^\{\(1\)\}is propagated backward to the final matrix\-multiplication node in the trace,VJP L1 Mm\.

#### 4\.2\.8VJP L1 Mm

In the primal pass, the input layer transformation ism\(1\)=t​\(W\(1\)\)⊤m^\{\(1\)\}=t\(W^\{\(1\)\}\)^\{\\top\}\. The nodeVJP L1 Mmreceives the incoming sensitivitym¯\(1\)∈ℝ1×3\\bar\{m\}^\{\(1\)\}\\in\\mathbb\{R\}^\{1\\times 3\}\. To complete the backward pass, the node evaluates the VJPs with respect to the temporal inputttand the first\-layer weightsW\(1\)W^\{\(1\)\}:

t¯\\displaystyle\\bar\{t\}=m¯\(1\)​W\(1\)\\displaystyle=\\bar\{m\}^\{\(1\)\}W^\{\(1\)\}\(23\)=\(−0\.3159−0\.19360\.6473\)​\(0\.2−0\.50\.8\)\\displaystyle=\\begin\{pmatrix\}\-0\.3159&\-0\.1936&0\.6473\\end\{pmatrix\}\\begin\{pmatrix\}0\.2\\\\ \-0\.5\\\\ 0\.8\\end\{pmatrix\}=−0\.06318\+0\.0968\+0\.51784=0\.5515,\\displaystyle=\-0\.06318\+0\.0968\+0\.51784=0\.5515,W¯\(1\)\\displaystyle\\bar\{W\}^\{\(1\)\}=\(m¯\(1\)\)⊤​t\\displaystyle=\(\\bar\{m\}^\{\(1\)\}\)^\{\\top\}t\(24\)=\(−0\.3159−0\.19360\.6473\)​\(0\.5\)=\(−0\.1579−0\.09680\.3236\)∈ℝ3×1\.\\displaystyle=\\begin\{pmatrix\}\-0\.3159\\\\ \-0\.1936\\\\ 0\.6473\\end\{pmatrix\}\(0\.5\)=\\begin\{pmatrix\}\-0\.1579\\\\ \-0\.0968\\\\ 0\.3236\\end\{pmatrix\}\\in\\mathbb\{R\}^\{3\\times 1\}\.
The scalart¯=0\.5515\\bar\{t\}=0\.5515is the terminal sensitivity at the input leaf, representing the numerical value of the total derivativey^′=∂y^/∂t\\hat\{y\}^\{\\prime\}=\\partial\\hat\{y\}/\\partial t\. In the context of a PINN, this value is not merely a gradient for optimization but is a primary component used to construct the differential equation residual\. Simultaneously, the column vectorW¯\(1\)\\bar\{W\}^\{\(1\)\}provides the terminal gradients for the first\-layer weights, ensuring the dimensions are consistent with the parameter matrixW\(1\)W^\{\(1\)\}in[Table˜3](https://arxiv.org/html/2607.13042#S3.T3)\.

##### A note on numerical precision\.

The valuet¯=0\.5515\\bar\{t\}=0\.5515calculated in this section exhibits a minor discrepancy compared to the value of0\.55130\.5513presented in[Section˜5](https://arxiv.org/html/2607.13042#S5)\. This is due to systematic rounding of intermediate sensitivities to four decimal places for didactic clarity\.

### 4\.3Summary

[Table˜6](https://arxiv.org/html/2607.13042#S4.T6)collects all adjoint values from the backward trace\. By initiating the pass with the seedy^¯=1\\bar\{\\hat\{y\}\}=1, we obtain the partial derivatives of the output with respect to every intermediate node and leaf parameter\.

Table 6:Reverse\-mode trace \(Graph 2 construction\) with seedy^¯=1\\bar\{\\hat\{y\}\}=1\. Each VJP node materializes a specific adjoint and buffers the necessary tensors \(both primal and adjoint\) required for higher\-order differentiation \(next backward pass\)\.VJP NodeProduces \(Operation = Value\)[grad\_fn](https://arxiv.org/html/2607.13042v1/grad_fn)Saved for backwardVJP L3 Addm¯\(3\)=y^¯⋅1=1∈ℝ\\bar\{m\}^\{\(3\)\}=\\bar\{\\hat\{y\}\}\\cdot 1=1\\in\\mathbb\{R\}[AddBackward0](https://arxiv.org/html/2607.13042v1/AddBackward0)\(none\)VJP L3 Mma¯\(2\)=m¯\(3\)​W\(3\)=\(0\.9,−0\.6,0\.3\)∈ℝ1×3\\bar\{a\}^\{\(2\)\}=\\bar\{m\}^\{\(3\)\}W^\{\(3\)\}=\(0\.9,\-0\.6,0\.3\)\\in\\mathbb\{R\}^\{1\\times 3\}[MmBackward0](https://arxiv.org/html/2607.13042v1/MmBackward0)m¯\(3\),W\(3\)\\bar\{m\}^\{\(3\)\},W^\{\(3\)\}VJP L2 Tanhz¯\(2\)=a¯\(2\)⊙ϕ′​\(z\(2\)\)=\(0\.83,−0\.58,0\.25\)∈ℝ1×3\\bar\{z\}^\{\(2\)\}=\\bar\{a\}^\{\(2\)\}\\odot\\phi^\{\\prime\}\(z^\{\(2\)\}\)=\(0\.83,\-0\.58,0\.25\)\\in\\mathbb\{R\}^\{1\\times 3\}[TanhBackward0](https://arxiv.org/html/2607.13042v1/TanhBackward0)a¯\(2\),a\(2\)\\bar\{a\}^\{\(2\)\},a^\{\(2\)\}VJP L2 Addm¯\(2\)=z¯\(2\)=\(0\.83,−0\.58,0\.25\)∈ℝ1×3\\bar\{m\}^\{\(2\)\}=\\bar\{z\}^\{\(2\)\}=\(0\.83,\-0\.58,0\.25\)\\in\\mathbb\{R\}^\{1\\times 3\}[AddBackward0](https://arxiv.org/html/2607.13042v1/AddBackward0)\(none\)VJP L2 Mma¯\(1\)=m¯\(2\)​W\(2\)=\(−0\.32,−0\.19,0\.67\)∈ℝ1×3\\bar\{a\}^\{\(1\)\}=\\bar\{m\}^\{\(2\)\}W^\{\(2\)\}=\(\-0\.32,\-0\.19,0\.67\)\\in\\mathbb\{R\}^\{1\\times 3\}[MmBackward0](https://arxiv.org/html/2607.13042v1/MmBackward0)m¯\(2\),W\(2\)\\bar\{m\}^\{\(2\)\},W^\{\(2\)\}VJP L1 Tanhz¯\(1\)=a¯\(1\)⊙ϕ′​\(z\(1\)\)=\(−0\.32,−0\.19,0\.65\)∈ℝ1×3\\bar\{z\}^\{\(1\)\}=\\bar\{a\}^\{\(1\)\}\\odot\\phi^\{\\prime\}\(z^\{\(1\)\}\)=\(\-0\.32,\-0\.19,0\.65\)\\in\\mathbb\{R\}^\{1\\times 3\}[TanhBackward0](https://arxiv.org/html/2607.13042v1/TanhBackward0)a¯\(1\),a\(1\)\\bar\{a\}^\{\(1\)\},a^\{\(1\)\}VJP L1 Addm¯\(1\)=z¯\(1\)=\(−0\.32,−0\.19,0\.65\)∈ℝ1×3\\bar\{m\}^\{\(1\)\}=\\bar\{z\}^\{\(1\)\}=\(\-0\.32,\-0\.19,0\.65\)\\in\\mathbb\{R\}^\{1\\times 3\}[AddBackward0](https://arxiv.org/html/2607.13042v1/AddBackward0)\(none\)VJP L1 Mmt¯=m¯\(1\)​W\(1\)=0\.5515∈ℝ\\bar\{t\}=\\bar\{m\}^\{\(1\)\}W^\{\(1\)\}=0\.5515\\in\\mathbb\{R\}[MmBackward0](https://arxiv.org/html/2607.13042v1/MmBackward0)m¯\(1\),W\(1\)\\bar\{m\}^\{\(1\)\},W^\{\(1\)\}[Figure˜7](https://arxiv.org/html/2607.13042#S4.F7)visualizes the resulting materialized adjoint graph, where the right\-to\-left orientation reflects the direction of the backward pass\. In this representation, dashed arrows indicate references back to the original primal nodes of Graph 1 required to evaluate the local Jacobians\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x8.png)Figure 7:The materialized adjoint graph \(Graph 2\)\. The structure follows the reverse\-mode sequence from top\-right to bottom\-left, capturing the chain of VJP operations and their dependencies on primal values for higher\-order differentiation\.

## 5The PINN Challenge: Higher\-Order Differentiation throughy^′\\hat\{y\}^\{\\prime\}

Unlike standard supervised learning, where gradients with respect to parametersθ\\thetasuffice, Physics\-Informed Neural Networks \(PINNs\) require differentiating the model’s output with respect to its input coordinates\. At a collocation pointtc=0\.5t\_\{c\}=0\.5, the ODE residualRRand its associated lossLRL\_\{R\}are defined as:

R=y^′\+y^,LR=R2⟹∂LR∂θ=2​R​\(∂y^∂θ\+∂y^′∂θ\)\.R=\\hat\{y\}^\{\\prime\}\+\\hat\{y\},\\qquad L\_\{R\}=R^\{2\}\\implies\\frac\{\\partial L\_\{R\}\}\{\\partial\\theta\}=2R\\left\(\\frac\{\\partial\\hat\{y\}\}\{\\partial\\theta\}\+\\frac\{\\partial\\hat\{y\}^\{\\prime\}\}\{\\partial\\theta\}\\right\)\.\(25\)The term∂y^′/∂θ\\partial\\hat\{y\}^\{\\prime\}/\\partial\\thetarequires differentiating the derivative computation itself—placing us exactly in the scenario discussed in[Section˜2\.5](https://arxiv.org/html/2607.13042#S2.SS5), wherettandθ\\thetanow take the roles of the independent variablesx1x\_\{1\}andx2x\_\{2\}from our earlier simplified example\.

In PyTorch, this initial backward AD is initiated by the command:

dy\_dt=torch\.autograd\.grad\(y\_hat,t,create\_graph=True\)\[0\]

Crucially, the flagcreate\_graph=Trueinstructs the engine to materialize the VJP operations themselves as new nodes in a secondary computational graph \(Graph 2\)\. This graph defines the functiony^′=∂y^/∂t\\hat\{y\}^\{\\prime\}=\\partial\\hat\{y\}/\\partial tin a differentiable form, physically linking the derivative back to the parametersθ\\thetathrough references to the primal weights and activations ofGraph 1\.

Omitting this flag leads to the same “Level 1” connectivity failure demonstrated in[Section˜2\.6](https://arxiv.org/html/2607.13042#S2.SS6):y^′\\hat\{y\}^\{\\prime\}remains numerically correct but becomes a constant leaf node, causing∂y^′/∂θ\\partial\\hat\{y\}^\{\\prime\}/\\partial\\thetato vanish\. As shown in[Table˜7](https://arxiv.org/html/2607.13042#S5.T7), this error is catastrophic\. For a specific weightW1,2\(2\)W^\{\(2\)\}\_\{1,2\}, the sensitivities∂y^/∂θ\\partial\\hat\{y\}/\\partial\\thetaand∂y^′/∂θ\\partial\\hat\{y\}^\{\\prime\}/\\partial\\thetahave opposite signs; when the latter is lost, the total gradient flips sign entirely\. Consequently, the optimizer moves the parameter in the opposite of the required direction, yet PyTorch raises no error or warning, as the operation remains mathematically valid but physically incomplete\.

Table 7:Impact ofcreate\_graphon the total PINN gradient for weightW1,2\(2\)W^\{\(2\)\}\_\{1,2\}\. Omission of the graph materialization disconnects the derivative sensitivity, causing a hidden error where the total gradient sign flips\.The actual computation of these sensitivities occurs during the traversal of this newly created structure\.[Figure˜8](https://arxiv.org/html/2607.13042#S5.F8)illustrates this “backward\-on\-backward” pass\. By seeding the terminal node ofGraph 2withd​y^′/d​y^′=1d\\hat\{y\}^\{\\prime\}/d\\hat\{y\}^\{\\prime\}=1, the engine propagates over\-adjoints through the VJP nodes\. The key mechanism is thecross\-link: the VJP nodes reference primal values \(such asz\(1\)z^\{\(1\)\}\) fromGraph 1to evaluate their own local Jacobians\. This jump between graphs allows the sensitivity ofy^′\\hat\{y\}^\{\\prime\}to reach the original parameter leaves, successfully yielding the mixed derivatives∂2y^/∂θ​∂t\\partial^\{2\}\\hat\{y\}/\\partial\\theta\\partial tand the higher\-order coordinate derivativey^′′\\hat\{y\}^\{\\prime\\prime\}\.

![Refer to caption](https://arxiv.org/html/2607.13042v1/x9.png)Figure 8:Pruned representation of the second\-order traversal\. The purple path shows the over\-adjoint flow through the VJP nodes of Graph 2\. The cross\-links between the Adjoint and Primal graphs provide the physical connectivity necessary to differentiate the gradient function with respect to the network parameters\.The numerical execution of this second backward pass is summarized in[Table˜8](https://arxiv.org/html/2607.13042#S5.T8)\. By initiating the traversal with the over\-adjoint seedy^′¯¯=1\\bar\{\\bar\{\\hat\{y\}^\{\\prime\}\}\}=1, we propagate the sensitivities through the VJP nodes\. Note how the over\-adjoint values change as they pass through the non\-linear Tanh VJP nodes, where the cross\-links toGraph 1primal activations \(z\(ℓ\)z^\{\(\\ell\)\}\) allow for the inclusion of the second\-orderϕ′′\\phi^\{\\prime\\prime\}terms\.

Table 8:Second\-order backward pass \(Traversal of Graph 2\)\. This pass propagates the over\-adjoint seed to compute the second\-order coordinate derivativey^′′\\hat\{y\}^\{\\prime\\prime\}and the parameter sensitivities∂y^′/∂θ\\partial\\hat\{y\}^\{\\prime\}/\\partial\\theta\.The final valuet¯¯=0\.2184\\bar\{\\bar\{t\}\}=0\.2184represents the second\-order temporal derivativey^′′\\hat\{y\}^\{\\prime\\prime\}\. Simultaneously, the intermediate over\-adjoints accumulated at the Mm nodes provide the sensitivities∂y^′/∂W\(ℓ\)\\partial\\hat\{y\}^\{\\prime\}/\\partial W^\{\(\\ell\)\}, completing the set of values required to evaluate the total PINN gradient previously shown in[Table˜7](https://arxiv.org/html/2607.13042#S5.T7)\.

### 5\.1The Complete PINN Training Step

Assembling the full loss from[Section˜3\.1](https://arxiv.org/html/2607.13042#S3.SS1): the total loss isL=LR\+λ⋅LI​CL=L\_\{R\}\+\\lambda\\cdot L\_\{IC\}withLR=R2L\_\{R\}=R^\{2\}at the collocation point andLI​C=\(y^​\(0\)−1\)2L\_\{IC\}=\(\\hat\{y\}\(0\)\-1\)^\{2\}at the initial condition, weighted byλ=10\\lambda=10\. The following code implements one complete training step:

Listing 7:One PINN training step with correct graph handling\. Parameter initialization follows Listing 1\.1importtorch

2importtorch\.nnasnn

3

4

5net=nn\.Sequential\(

6nn\.Linear\(1,3\),nn\.Tanh\(\),

7nn\.Linear\(3,3\),nn\.Tanh\(\),

8nn\.Linear\(3,1\),

9\)

10

11

12lam=10\.0

13

14

15tc=torch\.tensor\(\[\[0\.5\]\],requires\_grad=True\)

16y\_hat\_c=net\(tc\)

17

18

19dy\_dt=torch\.autograd\.grad\(

20y\_hat\_c,tc,create\_graph=True\)\[0\]

21

22

23R=dy\_dt\+y\_hat\_c

24loss\_R=R\*\*2

25

26

27t0=torch\.tensor\(\)

28y\_hat\_0=net\(t0\)

29loss\_IC=\(y\_hat\_0\-1\.0\)\*\*2

30

31

32loss=loss\_R\+lam\*loss\_IC

33loss\.backward\(\)

34

##### What[loss\.backward\(\)](https://arxiv.org/html/2607.13042v1/loss.backward())traverses\.

1. 1\.The loss nodes \(R2R^\{2\}, the weighted sum\), producing adjoint seed2​R=1\.45622R=1\.4562for the residual path and2​λ​\(y^​\(0\)−1\)=−22\.4202\\lambda\(\\hat\{y\}\(0\)\-1\)=\-22\.420for the initial condition path\.
2. 2\.The VJP graph \(Graph 2\), computing∂y^′/∂θk\\partial\\hat\{y\}^\{\\prime\}/\\partial\\theta\_\{k\}via the cross\-links to weights and activations\.
3. 3\.The forward graph \(Graph 1\), computing∂y^/∂θk\\partial\\hat\{y\}/\\partial\\theta\_\{k\}\.

The gradients from both paths are*accumulated*at each leaf, yielding the correct total:

param\.grad=2​R​\(∂y^∂θ\+∂y^′∂θ\)t=tc\+λ​2​\(y^​\(0\)−1\)​\(∂y^∂θ\)t=0\.\\texttt\{param\.grad\}=2R\\left\(\\frac\{\\partial\\hat\{y\}\}\{\\partial\\theta\}\+\\frac\{\\partial\\hat\{y\}^\{\\prime\}\}\{\\partial\\theta\}\\right\)\_\{t=t\_\{c\}\}\\\!\\\!\\\!\+\\;\\lambda\\,2\(\\hat\{y\}\(0\)\-1\)\\left\(\\frac\{\\partial\\hat\{y\}\}\{\\partial\\theta\}\\right\)\_\{t=0\}\.\(26\)

### 5\.2Key Points

1. 1\.The cost is roughly2×2\\times\.Graph 1 \(forward\) \+ Graph 2 \(VJP ofy^′\\hat\{y\}^\{\\prime\}\) \+ one backward pass through both\. Compared to standard training, PINNs pay approximately double in computation and memory\.
2. 2\.Higher\-order derivatives stack\.Forautograd\.grady^′′\\hat\{y\}^\{\\prime\\prime\}, apply[autograd\.grad](https://arxiv.org/html/2607.13042v1/autograd.grad)withcreate\_graph=Truetwice, building Graph 3 on top of Graph 2\. Each layer introduces one additional derivative ofϕ\\phi— akkth\-order ODE requiresϕ\(k\+1\)\\phi^\{\(k\+1\)\}in the parameter gradients\.
3. 3\.

## 6Practical Considerations

### 6\.1Thecreate\_graphPitfall

The comparison table in[Section˜5](https://arxiv.org/html/2607.13042#S5)showed that omittingcreate\_graph=Trueflips the sign of∂LR/∂W1,2\(2\)\\partial L\_\{R\}/\\partial W^\{\(2\)\}\_\{1,2\}from−0\.5619\-0\.5619to\+0\.0606\+0\.0606\. What makes this bug dangerous in practice is not the magnitude of the error — it is the absence of any signal that something is wrong:

- •PyTorch raises no error or warning\.
- •The loss may still decrease, because the∂y^/∂θ\\partial\\hat\{y\}/\\partial\\thetaterm alone provides some gradient signal — just the wrong one\.
- •The user may attribute poor accuracy to insufficient network capacity or suboptimal hyperparameters, never suspecting a one\-line bug\.

### 6\.2Gradient Accumulation

PyTorch accumulates gradients by default: each call to[loss\.backward\(\)](https://arxiv.org/html/2607.13042v1/loss.backward())*adds*to[param\.grad](https://arxiv.org/html/2607.13042v1/param.grad)rather than replacing it\. If gradients from the previous iteration are not cleared, the accumulated values are wrong — the optimizer sees a gradient computed at the current parameters plus a stale gradient computed at the previous parameters, a sum that has no mathematical meaning\.

The standard pattern is:

forepochinrange\(num\_epochs\):

optimizer\.zero\_grad\(\)

loss\.backward\(\)

optimizer\.step\(\)

The accumulation behavior exists by design: it allows gradient aggregation across mini\-batches when the full batch does not fit in memory\. Each mini\-batch contributes gradients from a disjoint subset of samples, and the sum reassembles the full\-batch gradient\. In PINNs, where the loss already combinesLRL\_\{R\}andLI​CL\_\{IC\}into a single scalar before calling[backward\(\)](https://arxiv.org/html/2607.13042v1/backward()), there is no mini\-batch splitting, and accumulation is simply a trap for the unaware\.

### 6\.3Memory

The graph\-on\-graph roughly doubles memory relative to standard training: Graph 1 stores intermediate activations; Graph 2 stores the adjoint intermediates and their cross\-links to Graph 1\. For higher\-order ODEs, each additional derivative adds another graph layer:

For the small networks typical of PINNs \(hundreds to thousands of parameters\), this overhead is manageable\. For larger architectures, gradient checkpointing trades computation for memory by discarding selected intermediate activations during the forward pass and recomputing them during the backward pass\. PyTorch provides this via[torch\.utils\.checkpoint](https://arxiv.org/html/2607.13042v1/torch.utils.checkpoint)\.

### 6\.4Numerical Precision

Theϕ′′\\phi^\{\\prime\\prime\}factor that appears in hidden\-layer gradients \([Section˜5](https://arxiv.org/html/2607.13042#S5)\) is:

ϕ′′​\(z\)=−2​tanh⁡\(z\)​\(1−tanh2⁡\(z\)\)\.\\phi^\{\\prime\\prime\}\(z\)=\-2\\tanh\(z\)\\bigl\(1\-\\tanh^\{2\}\(z\)\\bigr\)\.\(27\)At our operating point, the neurons are far from saturation\. For example, at Layer 2 neuron 1:ϕ′′​\(z1\(2\)\)=ϕ′′​\(0\.2837\)=−2​\(0\.2763\)​\(0\.9237\)=−0\.5104\\phi^\{\\prime\\prime\}\(z^\{\(2\)\}\_\{1\}\)=\\phi^\{\\prime\\prime\}\(0\.2837\)=\-2\(0\.2763\)\(0\.9237\)=\-0\.5104\. All factors areO​\(1\)O\(1\)and well resolved in any floating\-point format\.

The concern arises forsaturated neurons\(\|z\|≫1\|z\|\\gg 1\), wheretanh⁡\(z\)≈±1\\tanh\(z\)\\approx\\pm 1and1−tanh2⁡\(z\)≈01\-\\tanh^\{2\}\(z\)\\approx 0\. The productϕ′′\\phi^\{\\prime\\prime\}then involves multiplying a number near±2\\pm 2by a number near0\. In[float32](https://arxiv.org/html/2607.13042v1/float32), the small factor may lose significant digits, producing noisy or vanishing gradients\. The standard mitigation is double precision:

net=net\.double\(\)

t=torch\.tensor\(,dtype=torch\.float64,

requires\_grad=True\)

In practice, PINNs withtanh\\tanhactivation are routinely trained in[float64](https://arxiv.org/html/2607.13042v1/float64), especially when the loss involves second\-order or higher derivatives\.

### 6\.5Multiple Collocation Points

The running example uses a single collocation point for clarity\. WithNcN\_\{c\}collocation points,[autograd\.grad](https://arxiv.org/html/2607.13042v1/autograd.grad)requires a[grad\_outputs](https://arxiv.org/html/2607.13042v1/grad_outputs)argument to specify the adjoint seed for each point:

t\_col=torch\.linspace\(0,1,30,

requires\_grad=True\)\.unsqueeze\(1\)

y\_hat=net\(t\_col\)

dy\_dt=torch\.autograd\.grad\(

y\_hat,t\_col,

grad\_outputs=torch\.ones\_like\(y\_hat\),

create\_graph=True\)\[0\]

loss\_R=\(\(dy\_dt\+y\_hat\)\*\*2\)\.mean\(\)

Setting[grad\_outputs](https://arxiv.org/html/2607.13042v1/grad_outputs)to all ones means: compute∂yi/∂ti\\partial y\_\{i\}/\\partial t\_\{i\}for each collocation pointiiindependently\. Because the network is feedforward and eachyiy\_\{i\}depends only ontit\_\{i\}, the Jacobian∂𝐲/∂𝐭\\partial\\mathbf\{y\}/\\partial\\mathbf\{t\}is diagonal, and multiplying by the all\-ones vector recovers the diagonal entries — that is, the per\-point derivatives\.

## 7Conclusion

A PINN training step requires two levels of differentiation: first the physics derivativey^′=d​y^/d​t\\hat\{y\}^\{\\prime\}=d\\hat\{y\}/dt, then the parameter gradient∇θL\\nabla\_\{\\theta\}Lof a loss that depends ony^′\\hat\{y\}^\{\\prime\}\. The reader who has followed the worked example through this paper can now trace exactly how PyTorch handles both levels — and why a single missing flag silently breaks the second one\.

The paper exposed four mechanisms on a concrete 1\-3\-3\-1 MLP applied toy′\+y=0y^\{\\prime\}\+y=0,y​\(0\)=1y\(0\)=1:

1. 1\.Tangents and adjoints\([Section˜2](https://arxiv.org/html/2607.13042#S2)\): the chain rule evaluated forward \(one directional derivative per pass\) or backward \(all input derivatives in one pass\), using the same local Jacobians in both directions\.
2. 2\.The computational graph\([Section˜3](https://arxiv.org/html/2607.13042#S3)\): a trace of elementary operations, each saving exactly the tensors its VJP will need — nothing more\.
3. 3\.Reverse\-mode AD\([Section˜4](https://arxiv.org/html/2607.13042#S4)\): a single backward traversal that computes all 22 parameter gradients simultaneously, implicitly contracting with thePPsensitivities of\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\]without computing them explicitly\.
4. 4\.The graph\-on\-graph\([Section˜5](https://arxiv.org/html/2607.13042#S5)\): recording the VJP operations into a second graph so that[loss\.backward\(\)](https://arxiv.org/html/2607.13042v1/loss.backward())can traverse both, automatically producing the product rule andϕ′′\\phi^\{\\prime\\prime\}contributions that the companion paper derived by hand\. Omittingcreate\_graph=Truesilently drops the∂y^′/∂θ\\partial\\hat\{y\}^\{\\prime\}/\\partial\\thetaterm, producing wrong gradients without any error message\.

##### Relationship to\[[17](https://arxiv.org/html/2607.13042#bib.bib29)\]\.

The companion paper derives PINN gradients by hand using theP/QP/Qsensitivity framework — forward propagation of parameter perturbations, one parameter at a time \(22 passes\)\. This paper shows how PyTorch’s reverse\-mode engine arrives at the same numerical results by a single backward traversal through the computational graph\. The product rule that required explicitϕ′′\\phi^\{\\prime\\prime\}derivation in the companion paper emerges automatically from the chain rule through the graph\-on\-graph\. The two papers provide complementary views of the same computation: one algebraic, one algorithmic\.

##### Scope\.

We traced only a first\-order ODE on a small MLP withtanh\\tanhactivation\. The specific[grad\_fn](https://arxiv.org/html/2607.13042v1/grad_fn)nodes and saved tensors differ for other architectures and activation functions, but the three core mechanisms — the computational graph, the reverse\-mode traversal, and the graph\-on\-graph — do not\. Any differentiable network, any differential equation order, and any AD framework that supports graph recording will exhibit the same structure\.

##### Perspectives\.

- •Higher\-order and multi\-dimensional problems\.For PDEs requiring∇2y^\\nabla^\{2\}\\hat\{y\}, the graph\-on\-graph gains additional layers\. Tracing this structure for a model PDE would extend the present framework and expose the memory and precision challenges that arise when multiple graph layers accumulate\.
- •Mixed\-mode strategies\.Computingy^′\\hat\{y\}^\{\\prime\}by forward\-mode \(one pass, no graph\-on\-graph\) and∇θL\\nabla\_\{\\theta\}Lby reverse\-mode could reduce memory by eliminating Graph 2 entirely\. Both JAX \([jvp](https://arxiv.org/html/2607.13042v1/jvp)\+[grad](https://arxiv.org/html/2607.13042v1/grad)\) and PyTorch \([torch\.func\.jvp](https://arxiv.org/html/2607.13042v1/torch.func.jvp)\) support this natively\. Quantifying the memory–computation tradeoff on networks and PDEs of practical scale remains an open question\.

More broadly, the two\-level differentiation that PINNs require is not unique to them: any application that embeds a derivative in the objective — optimal control, differentiable simulation, neural ODEs — faces the same graph\-on\-graph structure\. The mechanisms traced here apply wherever a loss function depends on the derivative of a learned function\.

## Code and Data Availability

All results are generated from the network parameters in[Table˜3](https://arxiv.org/html/2607.13042#S3.T3)and the ODEy′\+y=0y^\{\\prime\}\+y=0,y​\(0\)=1y\(0\)=1\. A companion Jupyter notebook reproduces every calculation and verifies all intermediate values against PyTorch’s output\. The codebase is available at[https://github\.com/Tahimi/AD\-From\-Scratch\-PINN](https://github.com/Tahimi/AD-From-Scratch-PINN)and archived via Zenodo\[[16](https://arxiv.org/html/2607.13042#bib.bib31)\]\.

## Acknowledgments

During the development of this work, the author used AI\-based language models \(Claude, Anthropic, 2024–2025\) as assistive tools for writing, code development, and computational verification\. All conceptual and scientific decisions originated with the author\. All mathematical derivations, numerical values, and computational results were independently verified by the author, who assumes full responsibility for the integrity, accuracy, and originality of the submitted work\.

## References

- \[1\]\(2018\)Automatic differentiation in machine learning: a survey\.Journal of Machine Learning Research18\(153\),pp\. 1–43\.Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px3.p1.1)\.
- \[2\]J\. Blechschmidt and O\. G\. Ernst\(2021\)Three ways to solve partial differential equations with neural networks — a review\.GAMM\-Mitteilungen44\(2\),pp\. e202100006\.External Links:[Document](https://dx.doi.org/10.1002/gamm.202100006)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px3.p1.1)\.
- \[3\]S\. Cuomo, V\. S\. Di Cola, F\. Giampaolo, G\. Rozza, M\. Raissi, and F\. Piccialli\(2022\)Scientific machine learning through physics\-informed neural networks: where we are and what’s next\.Journal of Scientific Computing92\(3\),pp\. 88\.External Links:[Document](https://dx.doi.org/10.1007/s10915-022-01939-z)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[4\]M\. W\. M\. G\. Dissanayake and N\. Phan\-Thien\(1994\)Neural\-network\-based approximations for solving partial differential equations\.Communications in Numerical Methods in Engineering10\(3\),pp\. 195–201\.External Links:[Document](https://dx.doi.org/10.1002/cnm.1640100303)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[5\]W\. E and B\. Yu\(2018\)The deep ritz method: a deep learning\-based numerical algorithm for solving variational problems\.Communications in Mathematics and Statistics6\(1\),pp\. 1–12\.External Links:[Document](https://dx.doi.org/10.1007/s40304-018-0127-z)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[6\]I\. Goodfellow, Y\. Bengio, and A\. Courville\(2016\)Deep learning\.MIT Press\.External Links:[Link](http://www.deeplearningbook.org/)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px3.p1.1)\.
- \[7\]A\. Griewank and A\. Walther\(2008\)Evaluating derivatives: principles and techniques of algorithmic differentiation\.2nd edition,SIAM,Philadelphia\.External Links:[Document](https://dx.doi.org/10.1137/1.9780898717761)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px3.p1.1)\.
- \[8\]G\. E\. Karniadakis, I\. G\. Kevrekidis, L\. Lu, P\. Perdikaris, S\. Wang, and L\. Yang\(2021\)Physics\-informed machine learning\.Nature Reviews Physics3\(6\),pp\. 422–440\.External Links:[Document](https://dx.doi.org/10.1038/s42254-021-00314-5)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[9\]D\. Katsikis and A\. D\. Muradova\(2022\)A gentle introduction to physics\-informed neural networks, with applications in static rod and beam problems\.Journal of Advances in Applied & Computational Mathematics9,pp\. 103–127\.External Links:[Document](https://dx.doi.org/10.5757/JAACM.2022.9.103)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px3.p1.1)\.
- \[10\]I\. E\. Lagaris, A\. Likas, and D\. I\. Fotiadis\(1998\)Artificial neural networks for solving ordinary and partial differential equations\.IEEE Transactions on Neural Networks9\(5\),pp\. 987–1000\.External Links:[Document](https://dx.doi.org/10.1109/72.712178)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[11\]A\. Paszke, S\. Gross, F\. Massa, A\. Lerer, J\. Bradbury, G\. Chanan, T\. Killeen, Z\. Lin, N\. Gimelshein, L\. Antiga, A\. Desmaison, A\. Köpf, E\. Yang, Z\. DeVito, M\. Raison, A\. Tejani, S\. Chilamkurthy, B\. Steiner, L\. Fang, J\. Bai, and S\. Chintala\(2019\)PyTorch: an imperative style, high\-performance deep learning library\.InAdvances in Neural Information Processing Systems 32,pp\. 8024–8035\.Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[12\]M\. Raissi, P\. Perdikaris, and G\. E\. Karniadakis\(2019\)Physics\-informed neural networks: a deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations\.Journal of Computational Physics378,pp\. 686–707\.External Links:[Document](https://dx.doi.org/10.1016/j.jcp.2018.10.045)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[13\]M\. Raissi, A\. Yazdani, and G\. E\. Karniadakis\(2020\)Hidden fluid mechanics: learning velocity and pressure fields from flow visualizations\.Science367\(6481\),pp\. 1026–1030\.External Links:[Document](https://dx.doi.org/10.1126/science.aaw4741)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[14\]D\. E\. Rumelhart, G\. E\. Hinton, and R\. J\. Williams\(1986\)Learning representations by back\-propagating errors\.Nature323\(6088\),pp\. 533–536\.External Links:[Document](https://dx.doi.org/10.1038/323533a0)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px3.p1.1)\.
- \[15\]J\. Sirignano and K\. Spiliopoulos\(2018\)DGM: a deep learning algorithm for solving partial differential equations\.Journal of Computational Physics375,pp\. 1339–1364\.External Links:[Document](https://dx.doi.org/10.1016/j.jcp.2018.08.029)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.
- \[16\]A\. Tahimi\(2026\)Companion notebook for “automatic differentiation from scratch: how pytorch computes gradients in physics\-informed neural networks”\.Note:[https://github\.com/Tahimi/AD\-From\-Scratch\-PINN](https://github.com/Tahimi/AD-From-Scratch-PINN)Zenodo DOI to be assigned upon archivalCited by:[Code and Data Availability](https://arxiv.org/html/2607.13042#Sx1.p1.2)\.
- \[17\]A\. Tahimi\(2026\)Physics\-informed neural networks: a didactic derivation of the complete training cycle\.External Links:2604\.18481,[Link](https://arxiv.org/abs/2604.18481)Cited by:[4th item](https://arxiv.org/html/2607.13042#S1.I1.i4.p1.1),[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px1.p3.1),[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px4.p1.2),[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px4.p3.2),[item 3](https://arxiv.org/html/2607.13042#S7.I1.i3.p1.1),[§7](https://arxiv.org/html/2607.13042#S7.SS0.SSS0.Px1)\.
- \[18\]S\. Wang, S\. Sankaran, H\. Wang, and P\. Perdikaris\(2023\)An expert’s guide to training physics\-informed neural networks\.Computer Methods in Applied Mechanics and Engineering419,pp\. 116544\.External Links:[Document](https://dx.doi.org/10.1016/j.cma.2023.116544)Cited by:[§1](https://arxiv.org/html/2607.13042#S1.SS0.SSS0.Px2.p1.2)\.