Which direction costs what

A compiled JAX program will tell you what it thinks it costs. jax.jit(fn).lower(x).compile().cost_analysis() hands back a list with one dictionary in it, and flops is one of its keys. It is XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s own cost model rather than a hardware counter, so treat it as the compiler's arithmetic and not as a measurement of your machine. It repeats exactly across runs, which makes it a good instrument for comparing two programs.

Run it on the same function at three input widths and the shape of the cost model appears. The gradient sits at almost exactly twice the forward pass every time: 2.03, then 2.01, then 2.00. Widen the input by four and the ratio does not move, because reverse mode makes one backward sweep no matter how many inputs feed it.

jacfwd on the same function tells the opposite story. It costs 19 forward passes at 16 inputs, 68 at 64, and 260 at 256. That is the same tangent program run once per input direction, and the table below is what the chapter's cost argument looks like when you make the compiler count it.

import jax
import jax.numpy as jnp

def flops(fn, *args):
    return jax.jit(fn).lower(*args).compile().cost_analysis()[0]['flops']

n = 64
W = jnp.ones((n, n)) * 0.01

def f(x):
    return jnp.sum(jnp.tanh(x @ W))

x = jnp.ones(n)
print(flops(f, x), flops(jax.grad(f), x))
print(flops(jax.jacfwd(f), x), flops(jax.jacrev(f), x))

# 8255.0 16576.0
# 557056.0 16576.0

Turn the function around and the winner turns around with it. A map from 4 inputs to 256 outputs costs 13,600 flops through jacfwd and 788,736 through jacrev, a factor of 58 in favour of the direction that was hopeless a moment ago.

The rule behind both tables is one sentence about basis vectors. Forward mode runs once per input direction, reverse mode runs once per output direction, and the Jacobian needs every direction on one side or the other. Count inputs, count outputs, pick the smaller one.

Deep learning sits at one extreme of that rule and never moves: a billion parameters in, one loss out. The interesting cases are the ones that are not training steps, where a small parameter vector drives a large simulated trajectory and forward mode is the cheap direction by two orders of magnitude.

import jax
import jax.numpy as jnp

def flops(fn, *args):
    return jax.jit(fn).lower(*args).compile().cost_analysis()[0]['flops']

W = jnp.ones((4, 256)) * 0.01

def g(x):
    return jnp.tanh(x @ W)

x = jnp.ones(4)
print(flops(g, x))
print(flops(jax.jacfwd(g), x), flops(jax.jacrev(g), x))
print(jax.jacfwd(g)(x).shape)

# 2048.0
# 13600.0 788736.0
# (256, 4)

Open the wheel and the two builders are four lines each. jacfwd partially applies jvp at the point you passed, then vmaps it over _std_basis(dyn_args). jacrev calls _vjp, keeps the pullback, and vmaps that over _std_basis(y). Same shape of code, one built from the forward primitive and one from the reverse one.

_std_basis is where the flop ratios come from. It flattens the tree, counts the elements, and builds jnp.eye(ndim): an identity basis with one row per element of the thing it was handed. jacfwd gets the basis of the inputs, so it runs n times. jacrev gets the basis of the outputs, so it runs m times. The docs put the same fact in four words each, calling jacfwd column-by-column and jacrev row-by-row.

The out_axes=(None, -1) on the forward one is worth a second look. The primal output is shared across the batch rather than stacked, and the tangents stack on the last axis, which is what puts a Jacobian's columns where columns belong.

# jax/_src/api.py:581-582, inside jacfwd
      pushfwd: Callable = partial(_jvp, f_partial, dyn_args)
      y, jac = vmap(pushfwd, out_axes=(None, -1))(_std_basis(dyn_args))

# jax/_src/api.py:669 and 673, inside jacrev
      y, pullback = _vjp(f_partial, *dyn_args)
    jac = vmap(pullback)(_std_basis(y))

# jax/_src/api.py:763-769, the basis both of them batch over
def _std_basis(pytree):
  import jax.numpy as jnp
  leaves, _ = tree_flatten(pytree)
  ndim = sum(map(np.size, leaves))
  dtype = dtypes.result_type(*leaves)
  flat_basis = jnp.eye(ndim, dtype=dtype)
  return _unravel_array_into_pytree(pytree, 1, None, flat_basis)

The flop side of the bill favours reverse mode. The memory side is where it pays, and lesson one already showed where to look: the constvars of the linearized program. Count them for a block of tanh layers and the growth is exactly linear, 3 arrays at one layer and 17 at eight.

Print the shapes rather than the count and the arithmetic stops being abstract. One 128 by 128 weight matrix at 64 KiB, held once, plus two 32 by 128 activations per layer at 16 KiB each. Every layer you add costs 32 KiB at this batch size and nothing else, so a batch of 512 instead of 32 would cost 512 KiB per layer by the same rule.

Forward mode holds none of this, which is the other half of the trade the previous sections measured in flops. A tangent is consumed by the next equation; a residual has to survive until the backward sweep reaches it. LAB·J3 takes the next step from here, trading those bytes back for recomputation.

import jax
import jax.numpy as jnp

W = jnp.ones((128, 128)) * 0.01
x = jnp.ones((32, 128))

def residuals(depth):
    def block(x):
        for _ in range(depth):
            x = jnp.tanh(x @ W)
        return jnp.sum(x)
    _, lin = jax.linearize(block, x)
    avals = [v.aval for v in jax.make_jaxpr(lin)(x).jaxpr.constvars]
    return len(avals), sum(a.size * a.dtype.itemsize for a in avals)

for depth in (1, 2, 4, 8):
    print(depth, residuals(depth))

# 1 (3, 98304)
# 2 (5, 131072)
# 4 (9, 196608)
# 8 (17, 327680)