XAD documentation

Python tutorial

The following example for first-order adjoint mode illustrates how to use it:

import xad.adj_1st as xadj
# set independent variables
x0_ad = xadj.Real(1.0)
x1_ad = xadj.Real(1.5)
x2_ad = xadj.Real(1.3)
x3_ad = xadj.Real(1.2)
with xadj.Tape() as tape:
    # and register them
    tape.registerInput(x0_ad)
    tape.registerInput(x1_ad)
    tape.registerInput(x2_ad)
    tape.registerInput(x3_ad)
    # start recording derivatives
    tape.newRecording()
    # calculate the output
    y = x0_ad + x1_ad - x2_ad * x3_ad
    # register and seed adjoint of output
    tape.registerOutput(y)
    y.derivative = 1.0
    # compute all other adjoints
    tape.computeAdjoints()
    # output results
    print(f"y = {y}")
    print(f"first order derivatives:\n")
    print(f"dy/dx0 = {x0_ad.derivative}")
    print(f"dy/dx1 = {x1_ad.derivative}")
    print(f"dy/dx2 = {x2_ad.derivative}")
    print(f"dy/dx3 = {x3_ad.derivative}")

The Python bindings follow largely the same syntax and workflow as in C++.

Modules and Naming

Module Description Contents
xad The main module, which contain global functions and subpackages value, derivative, jacobian, hessian, get_include, get_cmake_dir, get_library_dir
xad.exceptions Contains all exceptions, with the same names as described in Exceptions e.g. NoTapeException
xad.math Mirrors Python's math module, with functions for XAD's active types. e.g. sin, exp
xad.cmath Mirrors Python's cmath module, for XAD's complex types. e.g. exp, log10
xad.fwd_1st First-order forward mode Real, Complex, Array
xad.adj_1st First-order adjoint mode and its tape Real, Complex, Array, Tape, CheckpointCallback, ScopedNestedRecording
xad.fwd_adj_2nd Second-order, forward over adjoint, and its tape Real, Complex, Tape, CheckpointCallback, ScopedNestedRecording
xad.fwd_fwd_2nd Second-order, forward over forward Real, Complex

Notes

  • The active type is called Real in all modes
  • In adjoint mode, a newly constructed Tape object is not automatically activated on construction. It can be activated using tape.activate() later, but we recommend using a with block as illustrated in the example above. Each thread has its own active tape.
  • The math functions in xad.math have been designed as a drop-in replacement for the standard Python math module. They not only support calls with XAD's active type, but also with regular float variables.
  • The x.getDerivative() and x.setDerivative() methods of active types are also available as the Python property x.derivative with both set and get functionality.
  • The x.getValue() method of active types is also available as the read-only property x.value
  • Use y.setDerivative(1.0) or the property setter y.derivative = 1.0 to seed and access derivatives.
  • Registering a variable as an input on a new tape always registers it afresh, also if it was an input on an earlier tape.

Complex Numbers

Complex has active real and imaginary parts, and xad.cmath mirrors Python's cmath (see samples/complex.py):

import xad.adj_1st as xadj
import xad.cmath
x, y = xadj.Real(0.5), xadj.Real(1.2)
with xadj.Tape() as tape:
    tape.registerInput(x)
    tape.registerInput(y)
    tape.newRecording()
    z = xadj.Complex(x, y)
    w = z * xad.cmath.exp(z)
    tape.registerOutput(w.real)
    w.real.derivative = 1.0
    tape.computeAdjoints()
    print(f"dRe(w)/dx = {x.derivative}, dRe(w)/dy = {y.derivative}")

Second-Order Derivatives

Second-order modes nest first-order types, as in C++: xad.fwd_adj_2nd (forward over adjoint) and xad.fwd_fwd_2nd. xad.value(x) and xad.derivative(x) return the inner first-order parts, so xad.derivative(xad.derivative(x)) is a second derivative (see samples/fwd_adj_2nd.py). For whole matrices, xad.jacobian and xad.hessian take a function of a list of Reals:

import xad
from xad.math import sin, cos
def f(x):
    return sin(x[0] * x[1]) - cos(x[1] * x[2]) - sin(x[2] * x[3]) - cos(x[3] * x[0])
print(xad.hessian(f, [1.0, 1.5, 1.3, 1.2]))

Checkpointing and External Functions

In adj_1st and fwd_adj_2nd, a subclass of CheckpointCallback replaces part of the recording; its computeAdjoint(tape) is called during computeAdjoints, as described in Checkpointing and External Functions for C++. The tape keeps inserted callbacks alive, and ScopedNestedRecording is a with block (see samples/checkpointing.py and samples/external_function.py):

import xad.adj_1st as xadj
class ExternalSum(xadj.CheckpointCallback):
    def compute_external(self, xs):
        tape = xadj.Tape.getActive()
        self.x_slots = [x.getSlot() for x in xs]
        y = xadj.Real(sum(x.value for x in xs))  # computed passively
        tape.registerOutput(y)
        self.y_slot = y.getSlot()
        tape.insertCallback(self)
        return y
    def computeAdjoint(self, tape):
        ybar = tape.getAndResetOutputAdjoint(self.y_slot)
        for s in self.x_slots:
            tape.incrementAdjoint(s, ybar)  # d(sum)/dx_i = 1

NumPy Arrays

xad.adj_1st.Array and xad.fwd_1st.Array hold n-dimensional arrays of active values in C++, so whole-array operations run as one C++ loop each instead of one Python call per element (see samples/numpy_arrays.py):

import numpy as np
import xad.adj_1st as xadj
with xadj.Tape() as tape:
    x = xadj.Array([0.1, 0.2, 0.3])
    tape.registerInputs(x)
    tape.newRecording()
    y = np.sum(np.arange(1.0, 4.0) * np.exp(x))
    tape.registerOutput(y)
    y.derivative = 1.0
    tape.computeAdjoints()
    print(x.derivatives)  # dy/dx as a float array

Supported are elementwise operators and ufuncs (arithmetic, comparisons, exp, log, sqrt, sin, ...) with broadcasting, sum/prod/mean with axis, @/dot, where, concatenate, and indexing, slicing and reshape as views. a[i, j] is the element Real itself, np.asarray(a) gives the values, and np.asarray(a, dtype=object) the elements. Unsupported NumPy functions raise TypeError rather than dropping the derivatives. Arrays need NumPy >= 1.25 at runtime; it is not a dependency of the package.

Building Extensions Against XAD

The package ships XAD as a shared library, with its C++ headers and a CMake package config. C++ extension modules linking it share the active tape with xad.adj_1st.Tape (per thread), as QuantLibAAD does:

find_package(XAD CONFIG REQUIRED)
target_link_libraries(my_ext PRIVATE XAD::xad)
cmake -B build -DXAD_DIR=$(python -c "import xad; print(xad.get_cmake_dir())")

xad.get_include() and xad.get_library_dir() give the header and library directories. Import xad before the extension so the library is loaded.

Questions or issues: GitHub Discussions · issues. Would you rather we AAD-enabled your code? Get a free speed-up estimate.