Autodiff as pushforward and pullback

Derivatives

Basic idea: jvp and vjp as pushforward and pullback on $\mathbb R^n$.

import jax
import jax.numpy as jnp

key = jax.random.key(42)

Example: $f: \mathbb R^n \to \mathbb R^2$.

[Note: python objects do have types but functions don’t declare or check them. Python does let us annotate functions, but the type would be jax.Array without specifying $n$ here. There are packages for annotating the latter.]

f = lambda x: jnp.array([
    jnp.sin(jnp.sum(x)),
    jnp.sum(x)]
)

# let's make it n = 3
x = jax.random.normal(key, (3,))

f(x)

Array([0.67023855, 0.7345302 ], dtype=float32)

Everything’s quite trivial here, but let’s call it $f: M \to N$ so we can write the input tangent space as $TM$ (instead of immediately identifying $T\mathbb R^n \cong \mathbb R^n$).

Forward ~ Pushforward

Then “forward mode” autodiff is

i.e. the “Jacobian * vector” product = jvp. Note that we specify all three objects simultaneously (the function, the input, and the tangent). That is because in forward mode all are computed alongside each other (autodiff = automated chain rule; in forward mode input and chain rule go in the same direction).

# note that this must be floats; [1, 0, 0] would give integer and jax would complain
v = jnp.array([1.0, 0.0, 0.0])

jax.jvp(f, (x,), (v,))

(Array([0.67023855, 0.7345302 ], dtype=float32),
 Array([0.7421457, 1.       ], dtype=float32))

Note $x$ and $v$ are wrapped in a tuple because that’s how jax deals with the issue that functions usually take multiple arguments. So for some other $g(x, y, z)$ the input space is $M = X \times Y \times Z$. For a single input, we still think of it as a “tuple” of one element.

Reverse ~ Pullback

For reverse mode, we instead propagate the cotangent through the chain rule, which goes the opposite direction from the input. Thus, vjp only needs us to give $f$ and $x$, and then returns a function to compute the pullback later.

Note the direction: vjp differentiates a map $M \to N$, but the function it hands back goes the other way, from cotangents on $N$ to cotangents on $M$. That is what makes it a pullback.

w = jnp.array([1.0, 1.0])

# note that we do *not* wrap (x,) here.
# by convention jax expects us to write it out, so it would be vjp(g, x, y, z) if
# we had a function with multiple input arguments
f_x, df_back = jax.vjp(f, x)
f_x  # consistency check; same output f(x) again

Array([0.67023855, 0.7345302 ], dtype=float32)

# now we can compute the "pullback" of w through f at x;
# returns a tuple, even for a single input,
# since again f is assumed to take a "product space of inputs" *always*
df_back(w)

(Array([1.7421458, 1.7421458, 1.7421458], dtype=float32),)

Note that we didn’t need to specify $x$ when evaluating df_back. That is because vjp already evaluated $f$ in the “forward” direction and “secretly” stored all intermediate values it computed inside of df_back for later use.

In terms of the python data types, for jax/numpy everything is just real numbers (or complex, see below). No distinction is made between $\mathbb R^n$, $T\mathbb R^n$ and $T^*\mathbb R^n$. In particular, note that both cotangents and tangents are represented as “row” vectors of shape (n,). To avoid confusion, it’s also worth noting that it works exactly the same for any array shape, so for “array” shaped inputs, inputs, tangents and cotangents are all (n, m) shaped numpy arrays.

Grad

jax.grad is just a convenience wrapper around the “pullback” jax.vjp.

The largest number of ML users just want to take the gradient of the loss function – a scalar function $l: \Theta \to L$ on the neural network parameters $\theta \in \Theta$. Of course $\Theta = \mathbb{R}^{\cdots}$ and $L=\mathbb R$. Since parameters are just real numbers, carrying the standard (flat) metric, we can immediately identify the gradient $\nabla l$ with the cotangent $1 \circ dl \in T^*\Theta$. To be extra verbose, the last expression includes the cotangent $1 \in T^*L$. This is exactly what the “convenience” function jax.grad implements.

l = lambda theta: jnp.sum(theta**2)

theta = jnp.array([[1.0, 2.0], [3.0, 4.0]])
grad_fn = jax.grad(l)

grad_fn(theta)  # 2 * theta

Array([[2., 4.],
       [6., 8.]], dtype=float32)

# jax.grad does a bit more book keeping, but the above is equivalent to:
def our_grad(f):
    # just like grad, return a *function* that computes the gradient

    def grad_fn(x):
        f_x, df_back = jax.vjp(f, x)
        # could also return f(x) -- exactly jax.value_and_grad
        return df_back(1.0)

    return grad_fn

# part of that book keeping: grad unwraps the one-element tuple, we don't
our_grad(l)(theta)

(Array([[2., 4.],
        [6., 8.]], dtype=float32),)

Jacobians

jax.jacfwd and jax.jacrev add no new capability either: they call jvp or vjp repeatedly, vectorized (vmap) over a canonical basis, to manifest the whole Jacobian as an array.

• jacfwd pushes each basis tangent $e_j \in TM$ forward, giving one column $df|_x(e_j)$ at a time.
• jacrev pulls each basis cotangent $dy^i \in T^*N$ back, giving one row $dy^i \circ df|_x$ at a time.

This is where the cost asymmetry comes from. For $f: \mathbb R^n \to \mathbb R^m$, forward mode needs $n$ passes (one per input dimension) and reverse mode needs $m$ passes (one per output dimension). Neither is universally better; reverse mode wins when there are many inputs and few outputs, which is another reason why grad (where $m = 1$) is naturally reverse mode.

def our_jacfwd(f, x):
    basis = jnp.eye(x.shape[0])                    # canonical basis of tangents
    push = lambda v: jax.jvp(f, (x,), (v,))[1]
    return jax.vmap(push)(basis).T                 # columns = pushed basis tangents

def our_jacrev(f, x):
    y, f_back = jax.vjp(f, x)
    basis = jnp.eye(y.shape[0])                    # canonical basis of cotangents
    pull = lambda w: f_back(w)[0]
    return jax.vmap(pull)(basis)                   # rows = pulled basis cotangents

# f: R^3 -> R^2 from above, so this costs 3 forward passes or 2 reverse passes
jnp.allclose(our_jacfwd(f, x), jax.jacfwd(f)(x)), jnp.allclose(our_jacrev(f, x), jax.jacrev(f)(x))

(Array(True, dtype=bool), Array(True, dtype=bool))

Complex numbers

In terms of vjp and jvp, jax treats $z \in \mathbb C$ identical to the vector [x, y] $\in \mathbb R^2$ such that $z = x + i y$. The function we differentiate through need not be holomorphic; it does what we would expect, thinking of $\mathbb C$ as a Riemannian manifold.

Concretely, decompose $z = x + iy$. Then $df = \frac{\partial f}{\partial x} \, dx + \frac{\partial f}{\partial y} \, dy$ or equivalently $df = \frac{\partial f}{\partial z} \, dz + \frac{\partial f}{\partial \bar z} \, d\bar{z}$. The mapping to how jax represents tangents and cotangents is straightforward, with the important note that the cotangent $v^\dagger$ is represented as $v$ (not as "$v^*$").

• $dx \cong$ 1.0
• $dy \cong$ 1.0j
• $dz \cong$ 1.0 + 1.0j
• $d\bar{z} \cong$ 1.0 - 1.0j
• $\partial_x \cong$ 1.0
• $\partial_y \cong$ 1.0j

Note that it might be tempting to expect to be able to read off $\partial_z f$ with a single jvp pass, by somehow specifying the tangent $\partial_z = (\partial_x - i \partial_y)/2$. However, the latter is a member of the complexified tangent space of which JAX knows nothing. Just looking at dimensions we see it can’t possibly work since the tangent space of $\mathbb C$, for jax, is two dimensional but $\partial_x, i\partial_x, \partial_y, i\partial_y$ are distinct objects in the complexified tangent space. In particular, if $\partial_y \cong 1.0j$ then $i\partial_y \cong -1.0$ but the latter is already the same as $-\partial_x$.

We thus have to implement the $\mathbb C$-linearity ourselves: $df(\partial_z) = df(\partial_x)/2 - i \, df(\partial_y)/2 \cong$ jvp(f, z, 1.0)/2 - 1j * jvp(f, z, 1.0j).

# consider complex conjugation
conj = lambda z: z.real - 1j * z.imag  #  = z.conj()

# note that z has to be complex if v is complex & vice versa
# removing one 0.0j yields a type error
z = 1.0 + 0.0j
v = 1.0 + 0.0j

out, out_tang = jax.jvp(conj, (z,), (v,))
out_tang  # directional derivative along the real axis gives 1

Array(1.+0.j, dtype=complex64, weak_type=True)

z = 1.0 + 0.0j
v = 0.0 + 1.0j

out, out_tang = jax.jvp(conj, (z,), (v,))
out_tang   # directional derivative along the imaginary axis gives -i

Array(0.-1.j, dtype=complex64, weak_type=True)

We can now implement the derivative for complexified tangent vectors (either taking complex coefficients of $c_x \partial_x + c_y \partial_y$ or equivalently $c_z \partial_z + c_{\bar z} \partial_{\bar z}$ as inputs).

def complexified_deriv(f, z, coeff_x, coeff_y):
    """Take derivative given complexified tangent vector."""
    # could also define equivalent function with coeff_z, coeff_zbar as inputs

    # evaluate real JVPs for \partial_x and \partial_y
    df_dx = jax.jvp(f, (z,), (1.0 + 0.0j,))[1]
    df_dy = jax.jvp(f, (z,), (0.0 + 1.0j,))[1]


    # combine by C-linearity
    return coeff_x * df_dx + coeff_y * df_dy

# partial_z (z) == 1, as it should
complexified_deriv(lambda z: z, 1.0 + 0j, 0.5, -0.5j)

(1+0j)

In particular we can now check if $f$ is holomorphic via the Cauchy-Riemann equation $\partial_{\bar z} f = 0$.

partial_z = lambda f, z: complexified_deriv(f, z, 0.5, -0.5j)
partial_zbar = lambda f, z: complexified_deriv(f, z, 0.5, 0.5j)

partial_zbar(jnp.sin, 1.0j)  # is holomorphic, so get 0

Array(0.+0.j, dtype=complex64, weak_type=True)

partial_zbar(jnp.conj, 1.0j)  # not holomorphic, don't get 0 (in fact is antiholomorphic, so get 1)

Array(1.+0.j, dtype=complex64, weak_type=True)

If we know a function is holomorphic, then Cauchy-Riemann $\partial_{\bar z} f = (\partial_x + i \partial_y) f / 2 = 0$ gives us $\partial_y f = i \, \partial_x f$. Thus, we could save on one of the jvp calls: df_dy = 1j * df_dx. We could thus in principle define two holomorphic derivatives, one “safe” one that works on any function, and one “unsafe” one that is cheaper but assumes f is holomorphic.

The analog story holds for vjp. We can again define a general version that works for the complexified cotangent space, and an “unsafe” one that assumes f is holomorphic.

# check the shortcut on a holomorphic function, before relying on it
holo = lambda z: z**3 + 2 * z
df_dx = jax.jvp(holo, (0.7 - 0.4j,), (1.0 + 0.0j,))[1]
df_dy = jax.jvp(holo, (0.7 - 0.4j,), (0.0 + 1.0j,))[1]

df_dy, 1j * df_dx  # equal, so one jvp would have sufficed

(Array(1.6800001+2.99j, dtype=complex64, weak_type=True),
 Array(1.6800001+2.99j, dtype=complex64, weak_type=True))

def complexified_cotangent_deriv(f, z, coeff_dx, coeff_dy):
    """ Evaluates the pull-back of a complexified cotangent."""

    _, vjp_fun = jax.vjp(f, z)

    # evaluate VJPs along standard real dual basis
    pullback_dx = vjp_fun(1.0 + 0.0j)[0]  # cotangent w = 1, real part pull-back
    pullback_dy = vjp_fun(0.0 + 1.0j)[0]  # cotangent w = i, imag part pull-back


    # combine by C-linearity
    return coeff_dx * pullback_dx + coeff_dy * pullback_dy

For scalar functions, we may of course again want to apply jax.grad. If the function outputs real numbers and the input is complex, everything works as expected. If the output is complex, however, it’s ambiguous which cotangent (in the two-dimensional real tangent space) to start with. For a holomorphic function $dy \circ df = i \, dx \circ df$ (where now $x, y$ decompose the output space $\mathbb C$), so we can just compute $dx \circ df$ and have all information. This is why jax.grad has an optional argument holomorphic and if set True it applies vjp starting with 1.0. For a multi-variable holomorphic function $f([z_1, z_2, z_3]) \in \mathbb C$, jax.grad(f, holomorphic=True) thus returns exactly the expected $z \mapsto [df/dz_1, df/dz_2, df/dz_3]$, assuming $f$ is holomorphic (else writing the complex derivative $df/dz_1$ is nonsense).

Splitting trick

There is one trick worth mentioning. If we know we have some non-holomorphic function $f$ of which we’ll want to take holomorphic and antiholomorphic derivatives, we can write it as a function of two arguments, $f(z, \bar{z})$. If we treat these two inputs as formally independent, and we promise that in the implementation of $f$ we only apply holomorphic functions in each of them, we can take the true (anti-) holomorphic derivative in a single backward/forward pass.

Exterior derivative operator

Given the number of quantities involved, there are many possible “book-keepings” that don’t change the math but change what we specify when and what we get out when. Some choices are suggested by efficiency. For example, we could compute the function value and the pushforward of a tangent vector in two separate passes, but if we know we want both that would be inefficient: the tangent propagation has to evaluate the forward chain of computations anyway. Nonetheless, different applications call for different computational structure, and the exterior derivative is a good place to see that play out: below we build it in forward mode, then try the obvious reverse-mode alternative and watch it fail to compose.

The convention here: a $k$-form is a function of a point and $k$ tangent vectors, omega(x, v_1, ..., v_k), returning its contraction. A $0$-form is then just a function of a point. For constant vector fields the exterior derivative is

where $\widehat{v_i}$ means that argument is omitted – so each term is one jvp of $\omega$ in the direction $v_i$, with the remaining tangents held fixed.

Note that the below is in many ways not the most efficient implementation (e.g. replacing the loop-reduce with a single vmaped call would be more efficient, keeping track if something was a total derivative would let us avoid computing $0$, …)

from functools import reduce

# helper: choose one out of n
def _pick_one(elements):
    parity = 1
    for i in range(len(elements)):
        yield parity, elements[i], (*elements[:i], *elements[i + 1:])
        parity *= -1

# take exterior derivative
def extd(fn):

    def d_fn(x, *tangents):
        return reduce(jnp.add, [
            jax.jvp(
                # want to take derivative only in x, so must hide ts dependence
                lambda _x: p * fn(_x, *ts),
                (x,), (t,)
            )[1]
            for p, t, ts in _pick_one(tangents)
        ])

    return d_fn

cube = lambda x: x**3
d_cube = extd(cube)  # d_cube(x, t) = 3 * x^2 * t

d_cube(3.0, 1/3)

Array(9., dtype=float32, weak_type=True)

Next, $d^2 = 0$. Note we have to leave $\mathbb R^1$ to test this properly: a $2$-form on a one-dimensional space vanishes for trivial reasons ($\Lambda^2 \mathbb R^1 = 0$), so it wouldn’t really test the cancellation. On $\mathbb R^2$ the two terms genuinely have to cancel.

h = lambda x: jnp.sin(x[0]) * x[1]**2  # a 0-form on R^2
ddh = extd(extd(h))

e1, e2 = jnp.array([1., 0.]), jnp.array([0., 1.])
ddh(jnp.array([0.3, -0.7]), e1, e2)

Array(0., dtype=float32)

As a slightly less trivial example, consider $g(x) = x^1\,dx^2 - x^2\,dx^1$. We might be tempted to immediately identify this with the (tangent) vector

and thus think we should implement $g$ as a map $\mathbb R^2 \to \mathbb R^2$. But that doesn’t quite fit the above, where we picked jvp and thus need to take tangent vectors as input. Thus we should implement $g$ as a map $\mathbb R^2 \to (T \mathbb R^2 \to \mathbb R)$ which implements the contraction $g(x) \cdot v$.

g = lambda x, v: v[1] * x[0] - v[0] * x[1]

g(jnp.array([0., 1.]), jnp.array([1., 0.]))  # -1

Array(-1., dtype=float32)

Now $g$ is not closed, $dg = 2 \, dx^1 \wedge dx^2$ so we can test extd on 1-forms.

dg = extd(g)

# g is not closed, so this is not always zero as ddh above was
dg(jnp.array([1., 1.]), e1, e2)

Array(2., dtype=float32)

# check orientation flip gives -1
dg(jnp.array([1., 1.]), e2, e1)

Array(-2., dtype=float32)

# collapse if tangents collinear
dg(jnp.array([1., 1.]), jnp.array([1., 1.]), jnp.array([2., 2.]))

Array(0., dtype=float32)

Why composition worked

Worth thinking about briefly, because the vjp case below is less straightforward. We represented a $k$-form as a function of a point and $k$ tangent vectors, contracted down to a number:

The output type is the input type with $k$ raised by one – the representation is closed under the operation. That is what lets us iterate: extd returns the same kind of object it consumes, so it can be fed back in. Note this is a statement about our chosen book-keeping, not about the mathematics (see below).

Reverse mode: the obvious thing to try

The forward version makes us pay one jvp per tangent we want to contract against. For a $0$-form that looks wasteful: $df|_x$ has $n$ components and reverse mode is supposed to hand us all of them in a single pass. So the obvious move is to write d with vjp instead.

Consider again $f: M \to N$. If $N = \mathbb R$ we can start from the “canonical” cotangent $1$; for general $N$ we have to say which cotangent on $N$ to project onto first.

def extd_rev(fn):

    def d_fn(x, cot=1.0):
        _, fn_back = jax.vjp(fn, x)
        (df_x,) = fn_back(cot)
        return df_x

    return d_fn

x0 = jnp.array([0.3, -0.7])

# same object as extd(h), but as cotangent components instead of a contraction:
# the forward version gives us one number per tangent we supply,
# the reverse version gives the whole covector in one pass.
extd_rev(h)(x0)

Array([ 0.46811488, -0.4137283 ], dtype=float32)

# can extract both with forward mode by projecting onto bases
(extd(h)(x0, e1), extd(h)(x0, e2))

(Array(0.46811488, dtype=float32), Array(-0.4137283, dtype=float32))

…but it only works once. If we try to iterate:

try:
    extd_rev(extd_rev(h))(x0)
except ValueError as err:
    print('ValueError:', err)

ValueError: unexpected JAX type (e.g. shape/dtype) for argument to VJP function: got float32[], but expected float32[2] because the corresponding output of the differentiated function had JAX type float32[2]

Two things go wrong.

The direction flips. jvp always runs input $\to$ output, so pushing forward again just continues in the same direction. But vjp returns a map into $T^*M$, cotangents on the input space. So omega = extd_rev(h) is a function $M \to T^*M$, and differentiating that in reverse mode needs a seed which pairs with a covector, i.e. a tangent, not a cotangent. One more level and it flips back. The seed type alternates $T^*, T, T^*,$ etc.

Also, note that the ValueError we did get is only the second point showing up by luck. In principle vjp could run another time, but only if we specify a projection cotangent, now of the shape of the space. That is because it makes no distinction between $\mathbb R^n$, $T\mathbb R^n$ and $T^*\mathbb R^n$, and so interprets the map generated by the first extd_rev as simply a map from $\mathbb R^n$ to $\mathbb R^n$. The bookkeeping is our responsibility.

We can of course still build the $2$-form in reverse mode, by doing the $n$ passes explicitly and antisymmetrizing. Which is precisely jacrev: vjp vectorized over a basis.

n = 3
basis = jnp.eye(n)
omega = lambda x: jnp.array([-x[1] * x[2], x[0], jnp.sin(x[0])])
x1 = jnp.array([0.3, -0.7, 1.1])

# one reverse pass per output component, with the flipped-type seed
slices = jnp.stack([extd_rev(omega)(x1, basis[j]) for j in range(n)])  # slices[j, i] = d_i omega_j
slices.T - slices  # (d omega)_{ij} = d_i omega_j - d_j omega_i

Array([[ 0.        ,  2.1       ,  0.25533652],
       [-2.1       ,  0.        ,  0.        ],
       [-0.25533652,  0.        ,  0.        ]], dtype=float32)

Closing the type again: components instead of contractions

The fix is not to patch extd_rev but to change the representation. Represent a $k$-form by its component array rather than by its contraction,

i.e. a function from a point to an antisymmetric array of shape (n,) * k. Now $d$ is one Jacobian plus an antisymmetrization,

and the output is again a component array, so the type closes and we can iterate. Note this version is agnostic about the mode: the Jacobian can come from jacfwd or jacrev and the answer is identical.

This representation does demand something the contraction version did not: the point has to be a shape-(n,) array, [0.1] rather than 0.1, even when $n = 1$. The reason is that the array axes are the form indices, so each derivative needs an axis of the point to hang its new index on, and a bare scalar has none. It is worth being strict about this rather than quietly calling atleast_1d on the Jacobian: that would make the one-dimensional case run, but on the second application there is still no input axis to distinguish from the form index, and the two collapse. The result is a plain second derivative instead of an antisymmetrization – for $x^3$ at $x = 0.1$ it returns $0.6$ where $d^2 = 0$ demands $0$. A wrong answer is worse than a failure, so we assert the shape instead. The contraction-based extd, by contrast, never indexes anything and so is indifferent to the encoding: d_cube(3.0, 1/3) above passed a bare float quite happily.

import math
from itertools import permutations

def _parity(perm):
    """Sign of a permutation, by counting inversions."""
    p, sgn = list(perm), 1
    for i in range(len(p)):
        for j in range(i + 1, len(p)):
            if p[i] > p[j]:
                sgn = -sgn
    return sgn

def antisym(arr):
    """Normalized antisymmetrization over all axes."""
    terms = [_parity(p) * jnp.transpose(arr, p) for p in permutations(range(arr.ndim))]
    return reduce(jnp.add, terms) / math.factorial(arr.ndim)

def extd_comp(omega, jac=jax.jacfwd):
    """Exterior derivative of a k-form given by components omega: M -> Lambda^k T*M.

    The point must be a shape-(n,) array -- [0.1], not 0.1 -- even for n = 1.
    """

    def d_omega(x):
        assert jnp.ndim(x) == 1, 'point must be a shape-(n,) array, not a bare scalar'
        # jac(omega)(x) has shape (n,) * k + (n,); move the input index to the front
        arr = jnp.moveaxis(jac(omega)(x), -1, 0)
        return arr.ndim * antisym(arr)  # arr.ndim == k + 1

    return d_omega

extd_comp(omega)(x1)

Array([[ 0.        ,  2.1       ,  0.25533652],
       [-2.1       ,  0.        ,  0.        ],
       [-0.25533652,  0.        ,  0.        ]], dtype=float32)

# cross-check the two representations against each other: contract omega by hand,
# hand it to the forward extd, and compare with one entry of the component version
omega_contracted = lambda x, v: jnp.dot(omega(x), v)

extd(omega_contracted)(x1, basis[0], basis[1]), extd_comp(omega)(x1)[0, 1]

(Array(2.1, dtype=float32), Array(2.1, dtype=float32))

# ...and now it composes, so d^2 = 0 is checkable on the manifest tensor
extd_comp(extd_comp(omega))(x1).max(), extd_comp(extd_comp(h))(x0).max()

(Array(0., dtype=float32), Array(0., dtype=float32))

# the mode genuinely does not matter for the result, only for the cost
jnp.allclose(extd_comp(omega, jax.jacfwd)(x1), extd_comp(omega, jax.jacrev)(x1))

Array(True, dtype=bool)