Jacobian-Vector Product of JAX AutoDiff Principle

[GoktugGuvercin]

Hello JAX Community;

While trying to use and understand jax.jvp() in forward accumulation of autodiff principle in JAX, I am confused about one point, so I wanted to ask it here:

As far as I understand, jax.jvp()function takes three major arguments which are the function to be differentiated, primals, and tangents. Primals refer to the values from input domain of target function, which will be also used for evaluation of jacobian f, whereas tangent vectors are the ones multiplied by the jacobian for the computation of JVP. At this point, there is a couple of details that we need to follow while passing these arguments:

• Pytree structure of primals need to be same as pytree structure of tangents
• The shape of each primal needs to be same as the shape of corresponding tangent

What I understand from these conditions is that there is a mutual relationship between primals and tangents: Each tangent vector needs to be of tangent space of domain of its corresponding primal. In other words, if the function takes 2 arguments (primals) in, for example, R^5 and R, we need to pass two tangent vectors, each of which has to be member tangent space of its corresponding primal domain, that is T(R^5) and T(R).

In this case, what I expected from jax.jvp() function is to compute 2 jacobians for those 2 primals, multiply these jacobians by 2 tangent vectors to map them into their corresponding cotangent spaces and return those 2 cotangent vectors. However, it seems like jax.jvp() sums those cotangent vectors and return that summation. What I cannot understand is why we don't get 2 cotangent vectors as return instead we get their sum if we pass 2 tangent vectors under the assumption of push-forwarding approach in differential geometry. Did I miss some detail or misunderstand some details ? Could you please clarify me about it.

I am adding an example code about this:

import jax
import jax.numpy as jnp
from jax import random

key = random.PRNGKey(0)
W_key, b_key, input_key, v1_key, v2_key = random.split(key, 5)

b = random.normal(b_key, (1,))
W = random.normal(W_key, (5,))
x = random.normal(input_key, (1,5))

tan_vec1 = random.normal(v1_key, W.shape)
tan_vec2 = random.normal(v2_key, b.shape)

single_neuron = lambda W, b: x @ W + b

y, cotan_vecs = jax.jvp(single_neuron, [W, b], [tan_vec1, tan_vec2])

jacobian_fun = jax.jacrev(single_neuron, argnums=(0, 1))
jacobian1, jacobian2 = jacobian_fun(W, b)

jvp1 = jacobian1 @ tan_vec1
jvp2 = jacobian2 @ tan_vec2

print("cotan vec 1 in R: ", jvp1)
print("cotan vec 2 in R: ", jvp2)
print("Sum of cotan vectors: ", jvp1 + jvp2)
print("What jax.jvp() returns: ", cotan_vecs)

[YouJiacheng]

You can think that JAX treats primals and tangents as a flattened vector.
i.e. in your R^5 and R example, JAX treats it as R^6.

[GoktugGuvercin]

I have just checked the source code, and I noticed that there is a small code block which supports your claim:

def _jvp(fun: lu.WrappedFun, primals, tangents, has_aux=False):
  """Variant of jvp() that takes an lu.WrappedFun."""
  if (not isinstance(primals, (tuple, list)) or
      not isinstance(tangents, (tuple, list))):
    raise TypeError("primal and tangent arguments to jax.jvp must be tuples or lists; "
                    f"found {type(primals).__name__} and {type(tangents).__name__}.")

  ps_flat, tree_def = tree_flatten(primals)
  ts_flat, tree_def_2 = tree_flatten(tangents)

As illustrated above, all primals and tangents are flattened, so each primal and corresponding tangent seem not to be regarded individually. Instead, they are like combined during flattening operation. Hence, it returns the summation, not cotangent vectors separately. Why was it designed like that ? I thought that each cotangent vectors, to which the tangent vectors are pushed forward, need to be accessed individually. I hope that @mattjj has very nice explanation about this.