Type demotion in `jax.jacobian`

Example:

from jax import jacobian
from jax.config import config
import jax.numpy as jnp
config.update('jax_enable_x64', True)

def f(x, y):
  return x * y

x = jnp.ones((), jnp.float32)
y = jnp.ones((), jnp.float64)

f(x, y).dtype, jacobian(f)(x, y).dtype

Gives

(dtype('float64'), dtype('float32'))

However, df/dx = y should have type jnp.float64.

(this seems to cause a type mismatch in https://github.com/google/neural-tangents/issues/112)

[froystig]

The observed behavior is correct, or at least consistent with our differentiation semantics. The cotangent of the first argument of f has type float32, since the first argument has type float32. The tangent of the output of f has type float16, since the output has type float16.

What’s perhaps unexpected is that jacobian is an alias for jacrev rather than jacfwd, as @jakevdp notes. I’d go further and offer that calling either function jacobian (via alias) is imprecise in this setting. Especially when a function has differing input/output types, it helps to take the more general view of its Jacobian (at a primal point) as a function.

Namely, for a fixed primal point x, the Jacobian J(x) is a (linear) function from input tangents to output tangents. Its transpose J'(x) is a function from output (co)tangents to input (co)tangents. In this view, jacfwd is the Jacobian map J(x) applied to the standard basis vectors of its domain, and jacrev is the Jacobian map’s transpose J'(x) applied to the standard basis vectors of its domain.

So the outputs of jacfwd and jacrev are elements of different vector spaces, typed accordingly. Neither object exactly captures the full notion of “the Jacobian”—that’s a function.

Perhaps we can improve documentation around jacobian, jacfwd, and jacrev to clarify their meaning.

[jakevdp]

Here’s a repro that’s slightly more to the point:

from jax import jacobian
import jax.numpy as jnp

def f(x):
  return x.astype('float16')

x = jnp.array(1.0)
print(f(x).dtype, jacobian(f)(x).dtype)
# float16 float32