Introducing THOAD, High Order Derivatives for PyTorch Graphs

mntsx — Tue, 02 Sep 2025 23:01:59 +0000

PRESENTING THE PACKAGE

I’m excited to share thoad (short for PyTorch High Order Automatic Differentiation), a Python only library that computes arbitrary order partial derivatives directly on a PyTorch computational graph. The package has been developed within a 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.

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 jax.jet or torch.func, thoad’s graph aware formulation enables:

Working with smaller pieced external derivatives
An optimization based on unifying independent dimensions (e.g. batch).

This delivers asymptotically better scaling with respect to order and batch size (respectively).

Additionally, we compute derivatives with a vectorial 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 PyTorch as its only dependency. Avoiding custom C++ or CUDA has a very positive impact on the long-term maintainability of the package.

The package can be installed from GitHub or PyPI:

GitHub: https://github.com/mntsx/thoad
PyPI: https://pypi.org/project/thoad/

thoad’s Hessian computation outperforms torch.autograd on GPU by 10–100× for practical networks, remaining close to the performance of the jax.jet implementation. See the notebook that reproduces the comparison: https://github.com/mntsx/thoad/blob/master/examples/benchmarks/benchmark\_vs\_torch\_autograd.ipynb.

thoad 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 backward. 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.

USING THE PACKAGE

thoad exposes two primary interfaces for computing high-order derivatives:

thoad.backward: a function-based interface that closely resembles torch.Tensor.backward. 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).
thoad.Controller: 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.

thoad.backward

The thoad.backward function computes high-order partial derivatives of a given output tensor and stores them in each leaf tensor’s .hgrad attribute.

Arguments:

tensor: A PyTorch tensor from which to start the backward pass. This tensor must require gradients and be part of a differentiable graph.
order: A positive integer specifying the maximum order of derivatives to compute.
gradient: A tensor with the same shape as tensor to seed the vector-Jacobian product (i.e., custom upstream gradient). If omitted, the default is used.
crossings: A boolean flag (default=False). If set to True, mixed partial derivatives (i.e., derivatives that involve more than one distinct leaf tensor) will be computed.
groups: An iterable of disjoint groups of leaf tensors. When crossings=False, only those mixed partials whose participating leaf tensors all lie within a single group will be calculated. If crossings=True and groups is provided, a ValueError will be raised (they are mutually exclusive).
keep_batch: A boolean flag (default=False) that controls how output dimensions are organized in the computed gradients.
When keep_batch=False:
The derivative preserves one first flattened "primal" axis, followed by each original partial shape, sorted in differentiation order. Concretelly:
- A single "primal" axis that contains every element of the graph output tensor (flattened into one dimension).
- A group of axes per derivative order, each matching the shape of the respective differentially targeted tensor.
For an N-th order derivative of a leaf tensor with input_numel elements and an output with output_numel elements, the deerivative shape is:
- Axis 1: indexes all output_numel outputs
- Axes 2…(sum(Nj)+1): each indexes all input_numel inputs
When keep_batch=True:
The derivative shape follows the same ordering as in the previous case, but includes a series of "independent dimensions" immediately after the "primal" axis.
- Axis 1 flattens all elements of the output tensor (size = output_numel).
- Axes 2...(k+i) correspond to dimensions shared by multiple input tensors and treated independently throughout the graph. These are dimensions that are only operated on element-wise (e.g. batch dimensions).
- Axes (k+i+1)...(k+i+sum(Nj)+1) each flatten all input_numel elements of the leaf tensor, one axis per derivative order.
keep_schwarz: A boolean flag (default=False). If True, symmetric (Schwarz) permutations are retained explicitly instead of being canonicalized/reduced—useful for debugging or inspecting non-reduced layouts.

Returns:

An instance of thoad.Controller wrapping the same tensor and graph.

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)

thoad.Controller

The Controller class wraps a tensor’s backward subgraph in a controller object, performing the same core high-order backward pass as thoad.backward while exposing advanced customization, inspection, and override capabilities.

Instantiation

Use the constructor to create a controller for any tensor requiring gradients:

controller = thoad.Controller(tensor=GO)  ## takes graph output tensor

tensor: A PyTorch Tensor with requires_grad=True and a non-None grad_fn.

Properties

.tensor → Tensor The output tensor underlying this controller. Setter: Replaces the tensor (after validation), rebuilds the internal computation graph, and invalidates any previously computed gradients.
.compatible → bool Indicates whether every backward function in the tensor’s subgraph has a supported high-order implementation. If False, some derivatives may fall back or be unavailable.
.index → Dict[Type[torch.autograd.Function], Type[ExtendedAutogradFunction]] A mapping from base PyTorch autograd.Function classes to thoad’s ExtendedAutogradFunction implementations. Setter: Validates and injects your custom high-order extensions.

Core Methods

.backward(order, gradient=None, crossings=False, groups=None, keep_batch=False, keep_schwarz=False) → None

Performs the high-order backward pass up to the specified derivative order, storing all computed partials in each leaf tensor’s .hgrad attribute.

order (int > 0): maximum derivative order.
gradient (Optional[Tensor]): custom upstream gradient with the same shape as controller.tensor.
crossings (bool, default False): If True, mixed partial derivatives across different leaf tensors will be computed.
groups (Optional[Iterable[Iterable[Tensor]]], default None): When crossings=False, restricts mixed partials to those whose leaf tensors all lie within a single group. If crossings=True and groups is provided, a ValueError is raised.
keep_batch (bool, default False): controls whether independent output axes are kept separate (batched) or merged (flattened) in stored/retrieved gradients.
keep_schwarz (bool, default False): if True, retains symmetric permutations explicitly (no Schwarz reduction).

.display_graph() → None

Prints a tree representation of the tensor’s backward subgraph. Supported nodes are shown normally; unsupported ones are annotated with (not supported).

.register_backward_hook(variables: Sequence[Tensor], hook: Callable) → None

Registers a user-provided hook to run during the backward pass whenever gradients for any of the specified leaf variables are computed.

variables (Sequence[Tensor]): Leaf tensors to monitor.
hook (Callable[[Tuple[Tensor, Tuple[Shape, ...], Tuple[Indep, ...]], dict[AutogradFunction, set[Tensor]]], Tuple[Tensor, Tuple[Shape, ...], Tuple[Indep, ...]]]): Receives the current (Tensor, shapes, indeps) plus contextual info, and must return the modified triple.

.require_grad_(variables: Sequence[Tensor]) → None

Marks the given leaf variables 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.

.fetch_hgrad(variables: Sequence[Tensor], keep_batch: bool = False, keep_schwarz: bool = False) → Tuple[Tensor, Tuple[Tuple[Shape, ...], Tuple[Indep, ...], VPerm]]

Retrieves the precomputed high-order partial corresponding to the ordered sequence of leaf variables.

variables (Sequence[Tensor]): the leaf tensors whose mixed partial you want.
keep_batch (bool, default False): if True, each independent output axis remains a separate batch dimension in the returned tensor; if False, independent axes are distributed/merged into derivative dimensions.
keep_schwarz (bool, default False): if True, returns derivatives retaining symmetric permutations explicitly.

Returns a pair:

Gradient tensor: the computed partial derivatives, shaped according to output and input dimensions (respecting keep_batch/keep_schwarz).
Metadata tuple
- Shapes (Tuple[Shape, ...]): the original shape of each leaf tensor.
- Indeps (Tuple[Indep, ...]): for each variable, indicates which output axes remained independent (batch) vs. which were merged into derivative axes.
- VPerm (Tuple[int, ...]): a permutation that maps the internal derivative layout to the requested variables order.

Use the combination of independent-dimension info and shapes to reshape or interpret the returned gradient tensor in your workflow.

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))

NOTE. A more detailed user guide with examples and feature walkthroughs is available in the notebook: https://github.com/mntsx/thoad/blob/master/examples/user_guide.ipynb

If you give it a try, I would love feedback on the API, corner cases, and models where you want better plug and play support.

DEV Community: mntsx

Introducing THOAD, High Order Derivatives for PyTorch Graphs

PRESENTING THE PACKAGE

USING THE PACKAGE

thoad.backward

thoad.Controller