Using XAD with Eigen

XAD's active types work as the scalar type of Eigen matrices and vectors. Eigen's arithmetic, reductions, decompositions and solvers then record on the tape like any other code, so derivatives flow through them without changes to Eigen. XAD is tested with Eigen 3.4.

Setup

Include XAD/StdCompatibility.hpp before XAD/XAD.hpp and the Eigen headers. It makes XAD's math functions available in namespace std, where Eigen looks for them, and specialises std::numeric_limits for the active types (see Headers and Namespaces).

#include <XAD/StdCompatibility.hpp>
#include <XAD/XAD.hpp>
#include <Eigen/Dense>

Matrices then hold the active type like any other scalar. Writing a function once as a template lets the same code run with double and with the active type:

template <typename T>
using Mat = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;
template <typename T>
using Vec = Eigen::Matrix<T, Eigen::Dynamic, 1>;
typedef xad::adj<double> mode;
typedef mode::active_type AD;
Mat<AD> A = A0.cast<AD>();   // A0 is a Mat<double>

Differentiating Through a Decomposition

This example, from the Eigen sample, computes the squared residual of a least-squares solution and its gradient with respect to the matrix A and the right-hand side b in one adjoint sweep:

template <typename T>
T computeLeastSquaresLoss(const Mat<T>& A, const Vec<T>& b)
{
    Vec<T> x = A.colPivHouseholderQr().solve(b);
    Vec<T> r = A * x - b;
    return r.squaredNorm();
}
mode::tape_type tape;
Mat<AD> A = A0.cast<AD>();
Vec<AD> b = b0.cast<AD>();
for (Eigen::Index j = 0; j < A.cols(); ++j)
    for (Eigen::Index i = 0; i < A.rows(); ++i) tape.registerInput(A(i, j));
for (Eigen::Index i = 0; i < b.size(); ++i) tape.registerInput(b(i));
tape.newRecording();
AD loss = computeLeastSquaresLoss(A, b);
tape.registerOutput(loss);
derivative(loss) = 1.0;
tape.computeAdjoints();
// derivative(A(i, j)) and derivative(b(i)) now hold dL/dA and dL/db

The tape records the whole QR factorisation and solve, so the adjoints account for how the solution x changes with A and b. The sample also computes log|det(M)| from the diagonal of a PartialPivLU factorisation and checks its gradient against the analytical result inverse(M)^T.

What Is Covered

XAD's tests and samples check derivatives through matrix products, inverse(), determinant(), trace(), norm(), PartialPivLU, ColPivHouseholderQR, Cholesky (llt()) and JacobiSVD, mostly in adjoint mode (inverse() also in forward mode), and with the direct mode types. Householder-based decompositions (QR, SVD, eigenvalue solvers) with newer Eigen releases, and stableNorm() / blueNorm(), compile with active types as well.

Comparisons and Tolerances

Active types have no default tolerance in Eigen. Eigen's comparisons at default precision, such as isApprox, isOnes, isIdentity and isUnitary, therefore compare exactly on active types. To compare with a tolerance, pass the precision explicitly, or compare the double values.

Sample

The Eigen sample contains the full least-squares and log-determinant code, with checks against finite differences. It is built when Eigen 3.4 or later is found by CMake, or downloaded with XAD_ENABLE_EIGEN_TESTS (see Tuning Behaviour and Performance):

cmake -S . -B build -DXAD_ENABLE_EIGEN_TESTS=ON
cmake --build build --target EigenSample
build/samples/Eigen/EigenSample