THOAD, Arbitrary Order Automatic Differentiation for PyTorch

username: mntsx
cooked: <h1><a name="p-475014-presenting-the-package-1" class="anchor" href="#p-475014-presenting-the-package-1"></a><strong>Presenting the package</strong></h1>
<p>I’m excited to finally release the first version of <strong>thoad</strong> (short for Py<strong>T</strong>orch <strong>H</strong>igh <strong>O</strong>rder <strong>A</strong>utomatic <strong>D</strong>ifferentiation), a Python only library that computes arbitrary order partial derivatives directly on a PyTorch computational graph. The package has been developed within a bachelor’s research project at Universidad Pontificia de Comillas - ICAI, and we are considering publishing a future academic article reviewing the mathematical details and the implementation design.</p>
<p>At its core, thoad takes a one output, many inputs view of the graph and pushes high order derivatives back to the leaf tensors. Although a 1→N problem can be rewritten as 1→1 by concatenating flattened inputs, as in functional approaches such as <code>jax.jet</code> or <code>functorch</code>, thoad’s graph aware formulation enables:</p>
<ul>
<li>
<p>Working with smaller <strong>pieced external derivatives</strong></p>
</li>
<li>
<p>An optimization based on <strong>unifying independent dimensions</strong> (especially batch).</p>
</li>
</ul>
<p>This delivers <strong>asymptotically better scaling</strong> with respect to order and batch size (respectively).</p>
<p>Additionally, we compute derivatives with a <em>vectorial</em> approach rather than component by component, which makes our pure PyTorch implementation possible. Consequently, the implementation stays at a high level, written entirely in Python and using <strong>PyTorch</strong> as its only dependency. Avoiding custom C++ or CUDA has a very positive impact on the long-term maintainability of the package.</p>
<p>The package is already available to be installed from <strong>GitHub</strong> or <strong>PyPI</strong>:</p>
<ul>
<li>GitHub: <a href="https://github.com/mntsx/thoad" rel="noopener nofollow ugc">https://github.com/mntsx/thoad</a></li>
<li>PyPI: <em>pip install thoad</em></li>
</ul>
<p>In our benchmarks, <strong>thoad outperforms</strong> <code>torch.autograd</code> <strong>for Hessian calculations even on CPU</strong>. See the repository <em>examples/benchmarks</em> to check the comparisons and run them in your own hardware.</p>
<p><strong>thoad</strong> is designed to align closely with PyTorch’s interface philosophy, so running the high order backward pass is practically indistinguishable from calling PyTorch’s own <code>backward</code>. When you need finer control, you can keep or reduce Schwarz symmetries, group variables to restrict mixed partials, and fetch the exact mixed derivative you need. Shapes and independence metadata are also exposed to keep interpretation straightforward.</p>
<h2><a name="p-475014-using-the-package-2" class="anchor" href="#p-475014-using-the-package-2"></a><strong>USING THE PACKAGE</strong></h2>
<p><strong>thoad</strong> exposes two primary interfaces for computing high-order derivatives:</p>
<ol>
<li><code>thoad.backward</code>: a function-based interface that closely resembles <code>torch.Tensor.backward</code>. It provides a quick way to compute high-order gradients without needing to manage an explicit controller object, but it offers only the core functionality (derivative computation and storage).</li>
<li><code>thoad.Controller</code>: a class-based interface that wraps the output tensor’s subgraph in a controller object. In addition to performing the same high-order backward pass, it gives access to advanced features such as fetching specific mixed partials, inspecting batch-dimension optimizations, overriding backward-function implementations, retaining intermediate partials, and registering custom hooks.</li>
</ol>
<h3><a name="p-475014-thoadbackward-3" class="anchor" href="#p-475014-thoadbackward-3"></a><strong>thoad.backward</strong></h3>
<p>The <code>thoad.backward</code> function computes high-order partial derivatives of a given output tensor and stores them in each leaf tensor’s <code>.hgrad</code> attribute.</p>
<p><strong>Arguments</strong>:</p>
<ul>
<li><code>tensor</code>: A PyTorch tensor from which to start the backward pass. This tensor must require gradients and be part of a differentiable graph.</li>
<li><code>order</code>: A positive integer specifying the maximum order of derivatives to compute.</li>
<li><code>gradient</code>: A tensor with the same shape as <code>tensor</code> to seed the vector-Jacobian product (i.e., custom upstream gradient). If omitted, the default is used.</li>
<li><code>crossings</code>: A boolean flag (default=<code>False</code>). If set to <code>True</code>, mixed partial derivatives (i.e., derivatives that involve more than one distinct leaf tensor) will be computed.</li>
<li><code>groups</code>: An iterable of disjoint groups of leaf tensors. When <code>crossings=False</code>, only those mixed partials whose participating leaf tensors all lie within a single group will be calculated. If <code>crossings=True</code> and <code>groups</code> is provided, a <em>ValueError</em> will be raised (they are mutually exclusive).</li>
<li><code>keep_batch</code>: A boolean flag (default=<code>False</code>) that controls how output dimensions are organized in the computed gradients.
<ul>
<li><strong>When</strong> <code>keep_batch=False</code>**:** Gradients are returned in a fully flattened form. Concretely, think of the gradient tensor as having:
<ul>
<li>A single “output” axis that lists every element of the original output tensor (flattened into one dimension).</li>
<li>One axis per derivative order, each listing every element of the corresponding input (also flattened).</li>
</ul>
</li>
<li>For an N-th order derivative of a leaf tensor with <code>input_numel</code> elements and an output with <code>output_numel</code> elements, the gradient shape is:
<ul>
<li><strong>Axis 1:</strong> indexes all <code>output_numel</code> outputs</li>
<li><strong>Axes 2…(N+1):</strong> each indexes all <code>input_numel</code> inputs</li>
</ul>
</li>
<li><strong>When</strong> <code>keep_batch=True</code>**:** Gradients preserve both a flattened “output” axis and each original output dimension before any input axes. You can visualize it as:
<ul>
<li><strong>Axis 1</strong> flattens all elements of the output tensor (size = <code>output_numel</code>).</li>
<li><strong>Axes 2…(k+1)</strong> correspond exactly to each dimension of the output tensor (if the output was shape <code>(d1, d2, ..., dk)</code>, these axes have sizes <code>d1</code>, <code>d2</code>, …, <code>dk</code>).</li>
<li><strong>Axes (k+2)…(k+N+1)</strong> each flatten all <code>input_numel</code> elements of the leaf tensor, one axis per derivative order.</li>
</ul>
</li>
<li>However, if a particular output axis does not influence the gradient for a given leaf, that axis is not expanded and instead becomes a size-1 dimension. This means only those output dimensions that actually affect a particular leaf’s gradient “spread” into the input axes; any untouched axes remain as 1, saving memory.</li>
</ul>
</li>
<li><code>keep_schwarz</code>: A boolean flag (default=<code>False</code>). If <code>True</code>, symmetric (Schwarz) permutations are retained explicitly instead of being canonicalized/reduced—useful for debugging or inspecting non-reduced layouts.</li>
</ul>
<p><strong>Returns</strong>:</p>
<ul>
<li>An instance of <code>thoad.Controller</code> wrapping the same tensor and graph.</li>
</ul>
<pre data-code-wrap="python"><code class="lang-python">import torch
import thoad
from torch.nn import functional as F

#### Normal PyTorch workflow
X = torch.rand(size=(10,15), requires_grad=True)
Y = torch.rand(size=(15,20), requires_grad=True)
Z = F.scaled_dot_product_attention(query=X, key=Y.T, value=Y.T)

#### Call thoad backward
order = 2
thoad.backward(tensor=Z, order=order)

#### Checks
## check derivative shapes
for o in range(1, 1 + order):
   assert X.hgrad[o - 1].shape == (Z.numel(), *(o * tuple(X.shape)))
   assert Y.hgrad[o - 1].shape == (Z.numel(), *(o * tuple(Y.shape)))
## check first derivatives (jacobians)
fn = lambda x, y: F.scaled_dot_product_attention(x, y.T, y.T)
J = torch.autograd.functional.jacobian(fn, (X, Y))
assert torch.allclose(J[0].flatten(), X.hgrad[0].flatten(), atol=1e-6)
assert torch.allclose(J[1].flatten(), Y.hgrad[0].flatten(), atol=1e-6)
## check second derivatives (hessians)
fn = lambda x, y: F.scaled_dot_product_attention(x, y.T, y.T).sum()
H = torch.autograd.functional.hessian(fn, (X, Y))
assert torch.allclose(H[0][0].flatten(), X.hgrad[1].sum(0).flatten(), atol=1e-6)
assert torch.allclose(H[1][1].flatten(), Y.hgrad[1].sum(0).flatten(), atol=1e-6)
</code></pre>
<h3><a name="p-475014-thoadcontroller-4" class="anchor" href="#p-475014-thoadcontroller-4"></a><strong>thoad.Controller</strong></h3>
<p>The <code>Controller</code> class wraps a tensor’s backward subgraph in a controller object, performing the same core high-order backward pass as <code>thoad.backward</code> while exposing advanced customization, inspection, and override capabilities.</p>
<p><strong>Instantiation</strong></p>
<p>Use the constructor to create a controller for any tensor requiring gradients:</p>
<pre><code>controller = thoad.Controller(tensor=GO)  ## takes graph output tensor
</code></pre>
<ul>
<li><code>tensor</code>: A PyTorch <code>Tensor</code> with <code>requires_grad=True</code> and a non-<code>None</code> <code>grad_fn</code>.</li>
</ul>
<p><strong>Properties</strong></p>
<ul>
<li><code>.tensor → Tensor</code> The output tensor underlying this controller. <strong>Setter</strong>: Replaces the tensor (after validation), rebuilds the internal computation graph, and invalidates any previously computed gradients.</li>
<li><code>.compatible → bool</code> Indicates whether every backward function in the tensor’s subgraph has a supported high-order implementation. If <code>False</code>, some derivatives may fall back or be unavailable.</li>
<li><code>.index → Dict[Type[torch.autograd.Function], Type[ExtendedAutogradFunction]]</code> A mapping from base PyTorch <code>autograd.Function</code> classes to thoad’s <code>ExtendedAutogradFunction</code> implementations. <strong>Setter</strong>: Validates and injects your custom high-order extensions.</li>
</ul>
<p><strong>Core Methods</strong></p>
<p><strong>.backward(order, gradient=None, crossings=False, groups=None, keep_batch=False, keep_schwarz=False) → None</strong></p>
<p>Performs the high-order backward pass up to the specified derivative <code>order</code>, storing all computed partials in each leaf tensor’s <code>.hgrad</code> attribute.</p>
<ul>
<li><code>order</code> (<code>int &gt; 0</code>): maximum derivative order.</li>
<li><code>gradient</code> (<code>Optional[Tensor]</code>): custom upstream gradient with the same shape as <code>controller.tensor</code>.</li>
<li><code>crossings</code> (<code>bool</code>, default <code>False</code>): If <code>True</code>, mixed partial derivatives across different leaf tensors will be computed.</li>
<li><code>groups</code> (<code>Optional[Iterable[Iterable[Tensor]]]</code>, default <code>None</code>): When <code>crossings=False</code>, restricts mixed partials to those whose leaf tensors all lie within a single group. If <code>crossings=True</code> and <code>groups</code> is provided, a <em>ValueError</em> is raised.</li>
<li><code>keep_batch</code> (<code>bool</code>, default <code>False</code>): controls whether independent output axes are kept separate (batched) or merged (flattened) in stored/retrieved gradients.</li>
<li><code>keep_schwarz</code> (<code>bool</code>, default <code>False</code>): if <code>True</code>, retains symmetric permutations explicitly (no Schwarz reduction).</li>
</ul>
<p><strong>.display_graph() → None</strong></p>
<p>Prints a tree representation of the tensor’s backward subgraph. Supported nodes are shown normally; unsupported ones are annotated with <code>(not supported)</code>.</p>
<p><strong>.register_backward_hook(variables: Sequence[Tensor], hook: Callable) → None</strong></p>
<p>Registers a user-provided <code>hook</code> to run during the backward pass whenever gradients for any of the specified leaf <code>variables</code> are computed.</p>
<ul>
<li><code>variables</code> (<code>Sequence[Tensor]</code>): Leaf tensors to monitor.</li>
<li><code>hook</code> (<code>Callable[[Tuple[Tensor, Tuple[Shape, ...], Tuple[Indep, ...]], dict[AutogradFunction, set[Tensor]]], Tuple[Tensor, Tuple[Shape, ...], Tuple[Indep, ...]]]</code>): Receives the current <code>(Tensor, shapes, indeps)</code> plus contextual info, and must return the modified triple.</li>
</ul>
<p><strong>.require_grad_(variables: Sequence[Tensor]) → None</strong></p>
<p>Marks the given leaf <code>variables</code> so that all intermediate partials involving them are retained, even if not required for the final requested gradients. Useful for inspecting or re-using higher-order intermediates.</p>
<p><strong>.fetch_hgrad(variables: Sequence[Tensor], keep_batch: bool = False, keep_schwarz: bool = False) → Tuple[Tensor, Tuple[Tuple[Shape, …], Tuple[Indep, …], VPerm]]</strong></p>
<p>Retrieves the precomputed high-order partial corresponding to the ordered sequence of leaf <code>variables</code>.</p>
<ul>
<li><code>variables</code> (<code>Sequence[Tensor]</code>): the leaf tensors whose mixed partial you want.</li>
<li><code>keep_batch</code> (<code>bool</code>, default <code>False</code>): if <code>True</code>, each independent output axis remains a separate batch dimension in the returned tensor; if <code>False</code>, independent axes are distributed/merged into derivative dimensions.</li>
<li><code>keep_schwarz</code> (<code>bool</code>, default <code>False</code>): if <code>True</code>, returns derivatives retaining symmetric permutations explicitly.</li>
</ul>
<p>Returns a pair:</p>
<ol>
<li><strong>Gradient tensor</strong>: the computed partial derivatives, shaped according to output and input dimensions (respecting <code>keep_batch</code>/<code>keep_schwarz</code>).</li>
<li><strong>Metadata tuple</strong>
<ul>
<li><strong>Shapes</strong> (<code>Tuple[Shape, ...]</code>): the original shape of each leaf tensor.</li>
<li><strong>Indeps</strong> (<code>Tuple[Indep, ...]</code>): for each variable, indicates which output axes remained independent (batch) vs. which were merged into derivative axes.</li>
<li><strong>VPerm</strong> (<code>Tuple[int, ...]</code>): a permutation that maps the internal derivative layout to the requested <code>variables</code> order.</li>
</ul>
</li>
</ol>
<p>Use the combination of independent-dimension info and shapes to reshape or interpret the returned gradient tensor in your workflow.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import thoad
from torch.nn import functional as F

#### Normal PyTorch workflow
X = torch.rand(size=(10,15), requires_grad=True)
Y = torch.rand(size=(15,20), requires_grad=True)
Z = F.scaled_dot_product_attention(query=X, key=Y.T, value=Y.T)

#### Instantiate thoad controller and call backward
order = 2
controller = thoad.Controller(tensor=Z)
controller.backward(order=order, crossings=True)

#### Fetch Partial Derivatives
## fetch T0 and T1 2nd order derivatives
partial_XX, _ = controller.fetch_hgrad(variables=(X, X))
partial_YY, _ = controller.fetch_hgrad(variables=(Y, Y))
assert torch.allclose(partial_XX, X.hgrad[1])
assert torch.allclose(partial_YY, Y.hgrad[1])
## fetch cross derivatives
partial_XY, _ = controller.fetch_hgrad(variables=(X, Y))
partial_YX, _ = controller.fetch_hgrad(variables=(Y, X))
</code></pre>
<blockquote>
<p>NOTE. A more detailed user guide with examples and feature walkthroughs is available in the notebook: <a href="https://github.com/mntsx/thoad/blob/master/examples/user_guide.ipynb" rel="noopener nofollow ugc">https://github.com/mntsx/thoad/blob/master/examples/user_guide.ipynb</a></p>
</blockquote>
<p>If you give it a try, I would love feedback on the API.</p>