JVP and VJP for tensor-valued function

[exenGT]

Hi,

How can I calculate the JVP and VJP for tensor-valued function, like the following:

Here, A is a tensor, x can be a vector or a tensor.

Thanks,
JW

[YouJiacheng]

jax.jvp and jax.vjp is compatible with tesnor-valued functions. No extra effort is needed.