Can I get JVP and VJP in one pass?

[qsh-zh]

Assume we have a function fn mapping from a to b, JVP is a mapping (a, T a) to (b, T b), while from one pass VJP we can get b, f_vjp.
Do we have JVP_and_JVP such that mapping (a, T a) to (b, T b, f_vjp)? I have a naive implementation

def naive_jvp_and_jvp(fn, a, T_a):
  b, T_b = jax.jvp(fn, (a, ), (T_a,))
  b, f_jvp = jax.vjp(fn, a)
  return b, T_b, f_jvp

However, such implementation needs to calculate b twice, how can I save the calculation and reach a more efficient implementation?

[soraros]

I think they'd be deduped after JIT.

[qsh-zh]

how can we verify that it is deduped?

[YouJiacheng]

You can verify that it is deduped by:

jax.jit(lambda x: x * x).lower(2.0).compile().compiler_ir()[0].to_string()

Alternatively you can:

def jvp_and_vjp(fn, a, T_a):
  b, jvp_f = jax.linearize(fn, a)
  T_b = jvp_f(T_a)
  vjp_f = jax.linear_transpose(jvp_f, T_b) # only use T_b shape and dtype
  return b, T_b, vjp_f