Compute the jacobian of `f` at `primals` multiplied by `tangents`.
(f, primals, tangents)
| 48 | |
| 49 | # TODO(allenl): Move this somewhere useful once forward gradients are stable. |
| 50 | def _jvp(f, primals, tangents): |
| 51 | """Compute the jacobian of `f` at `primals` multiplied by `tangents`.""" |
| 52 | with forwardprop.ForwardGradientAccumulator() as acc: |
| 53 | acc.watch(primals, tangents) |
| 54 | primals_out = f(*primals) |
| 55 | return primals_out, acc.jvp(primals_out) |
| 56 | |
| 57 | |
| 58 | def _jacfwd(f, primals): |