How to most efficiently calculate jvp when you want to calculate two tangents at the same primal?

[adam-hartshorne]

If I have a primal x, a variable z (that we don't want to calculate tangents over), and two tangents t1 and t2, where I want to calculate the primal_out and the tangent_out for both t1 and t2 at x.

I am aware of jax.linearize approach i.e

def test_fun(x, z):
    .....

y, f_jvp = jax.linearize(test_fun, x, z)
t1_out = f_jvp(t1, 0.0)
t2_out = f_jvp(t2, 0.0)

is this really the most efficient way, or with such a small number of tangents is there a better way that doesn't require creating f_jvp and two calls to it?

[soraros]

You could use a vmapp-ed f_jvp. Noticeably, this trick is also used to implement jacfwd and jacrev.

from functools import partial

import jax.numpy as jnp
from jax import jvp, linearize, vmap

ps = (jnp.arange(3.), 1.)

t1 = jnp.array([0, 0, 1.])
t2 = jnp.array([0, 1., 0])
ts = (jnp.stack([t1, t2]), jnp.stack([0., 0.]))

def f(x, z):
  return x + jnp.sin(z)

def f_jvps_manual(ps, ts):
  _, f_jvp = linearize(f, *ps)
  return tuple(f_jvp(*x) for x in zip(*ts))

def f_jvps_vmap(ps, ts):
  _, ts_out = vmap(partial(jvp, f, ps))(ts)
  return tuple(ts_out)

print(f_jvps_manual(ps, ts))
print(f_jvps_vmap(ps, ts))

[mattjj]

+1 to what @soraros wrote!

The jax.linearize approach and the vmapped-jax.jvp approach perform the same number of FLOPs. But batching things together is usually more efficient, because e.g. we get parallelism within operations.

The jax.linearize approach may be more useful if not all the vectors are available at the same time, e.g. in Krylov subspace methods.

[adam-hartshorne]

I also require the primal_out, "y", from the function call i.e. y, f_jvp = linearize(f, *ps)

Using the vmap approach defined above, I would have to capture 2 copies of y, which if y is large would be memory inefficient?

[soraros]

@adam-hartshorne If you have concerns for that, I suppose you could vmap the linearised function f.