The rule under grad

jax.jvp takes three things and returns two. The function, a tuple of primal arguments, a tuple of tangents matching them one for one, and back come the value and the directional derivative at that value. No Jacobian is built anywhere in that call, which is why it costs a forward pass rather than a Jacobian's worth of them.

The run below checks the tangent against a Jacobian built the slow way. jax.jacrev(f)(x) @ v assembles the full 2-by-2 matrix and multiplies; jax.jvp never assembles anything. On this input the two agree bit for bit, which is the difference between a derivative rule composed exactly and a finite difference that would agree to a few digits at best.

One shape detail decides how you call it. The primals and tangents are tuples over the function's positional arguments, not over the array's elements, so a one-argument function takes (x,) and (v,) with the commas that make them tuples. Forgetting a comma is the most common way this call fails.

import jax
import jax.numpy as jnp

def f(x):
    return jnp.stack([jnp.sin(x[0]) * x[1], x[0] ** 2 + x[1]])

x = jnp.array([0.5, 2.0])
v = jnp.array([1.0, 0.0])

y, dy = jax.jvp(f, (x,), (v,))
print(y, dy)
print(jnp.array_equal(dy, jax.jacrev(f)(x) @ v))

# [0.9588511 2.25     ] [1.7551651 1.       ]
# True

Trace sin(x) * x and you get two equations. Trace its jvp and you get seven, and every one of them is either a primal equation you already had or a tangent equation sitting next to it. cos a appears because the sine rule needs it. The two mul equations feeding add_any are the product rule, one term per factor, added together at the end.

Read the output line of the jvp jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → and the structure is plain: in (f, i), the primal result and the tangent result, produced by one program in one pass. There is no second program, no tape, and nothing held back for later. A tangent is consumed by the next equation the moment it is produced.

That is why forward mode's memory cost is the same as the forward pass's. It also explains why the cost scales with the number of directions you ask about: one tangent vector per call, and the tangent equations run once for each.

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x) * x

x = jnp.arange(3.0)
print(jax.make_jaxpr(f)(x))
print(jax.make_jaxpr(lambda x, v: jax.jvp(f, (x,), (v,)))(x, jnp.ones(3)))

# { lambda ; a:f32[3]. let b:f32[3] = sin a; c:f32[3] = mul b a in (c,) }
# { lambda ; a:f32[3] b:f32[3]. let
#     c:f32[3] = sin a
#     d:f32[3] = cos a
#     e:f32[3] = mul b d
#     f:f32[3] = mul c a
#     g:f32[3] = mul e a
#     h:f32[3] = mul c b
#     i:f32[3] = add_any g h
#   in (f, i) }

jax.linearize(f, x) runs the primal half once and gives you back the value plus a function that is linear in the tangent. The docs describe it in one line as producing a linear approximation using jvp() and partial eval, and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → of what comes back shows exactly which half got evaluated.

Look at where the semicolon falls. Three arrays sit before it, as constvars, and only one variable sits after it, the tangent. Those three constants are the residuals: the values from the forward pass that the linear program still needs. Print them and they are cos(x), x, and sin(x), in that order, which are precisely the three factors the product rule and the sine rule asked for.

This is the memory that reverse mode is famous for spending, made countable. It is not a mystery buffer inside the framework; it is a list of arrays you can print, and the next lesson counts them for a deep block.

The residuals are not hidden. They are the constvars of the linearized program.

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sum(jnp.sin(x) * x)

x = jnp.arange(3.0)
y, f_lin = jax.linearize(f, x)
lin = jax.make_jaxpr(f_lin)(jnp.ones(3))
print(lin)
print([c.tolist() for c in lin.consts])

# { lambda a:f32[3] b:f32[3] c:f32[3]; d:f32[3]. let
#     e:f32[3] = mul d a
#     f:f32[3] = mul e b
#     g:f32[3] = mul c d
#     h:f32[3] = add_any f g
#     i:f32[] = reduce_sum[axes=(0,)] h
#   in (i,) }
# [[1.0, 0.5403022766113281, -0.416146844625473], [0.0, 1.0, 2.0], [0.0, 0.8414709568023682, 0.9092974066734314]]

A linear function has a transpose, and JAX will build it for you. Hand jax.linear_transpose the linear part from the previous section and the point it was linearized at, feed the transposed function a cotangent of 1.0, and what comes back is the gradient, element for element the same array jax.grad returns.

So a primitive never registers a backward rule. It registers a forward one, JAX linearizes the composition, and transposition turns that linear program around. Reverse mode is a consequence of two transformations rather than a second differentiation algorithm sitting beside the first.

The Pallas arc reaches this same decomposition from the other end, quoting the design document on why a kernel with overlapping parallel reads transposes into something slow. Same three steps, read there as a limit on what can be compiled; read here as where your gradient comes from.

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sum(jnp.sin(x) * x)

x = jnp.arange(3.0)
_, f_lin = jax.linearize(f, x)
print(jax.linear_transpose(f_lin, x)(1.0))
print(jax.grad(f)(x))

# (Array([0.        , 1.3817732 , 0.07700372], dtype=float32),)
# [0.         1.3817732  0.07700372]

Tangents live in a different space from primals, and JAX enforces the match rather than casting quietly. Hand a float16 tangent to a float32 primal and the call raises with a message that names both dtypes and the one it expected. Nothing is promoted for you, because a promotion here would change the numerics of a derivative you asked to be exact.

Integer primals get the interesting answer. The tangent space of an integer array has no directions in it, so JAX gives it a dtype called float0 whose itemsize really is zero bytes. Differentiate through an integer with allow_int=True and what comes back is a float0 array of the right shape: a stated absence, not a buffer of zeros you might mistake for a gradient.

import numpy as np
import jax
import jax.numpy as jnp
from jax import dtypes

try:
    jax.jvp(lambda t: t * 2.0, (jnp.arange(3.0),), (jnp.ones(3, dtype=jnp.float16),))
except TypeError as err:
    print(err)

n = jnp.arange(3)
print(jax.jvp(lambda t: t * 2, (n,), (np.zeros(3, dtype=dtypes.float0),))[1].dtype)
print(np.dtype(dtypes.float0).itemsize)
print(jax.grad(lambda v: jnp.sum(v * 2.0), allow_int=True)(n).dtype)

# primal and tangent arguments to jax.jvp do not match; dtypes must be equal, or in
# case of int/bool primal dtype the tangent dtype must be float0.Got primal dtype
# float32 and so expected tangent dtype float32, but got tangent dtype float16 instead.
# [('float0', 'V')]
# 0
# [('float0', 'V')]