Hi,

We actually do that by doing vJ with a batch of v so you won’t be able to re-use it here.
That would be “element-wise” jacobian vector product. There are two cases here:
a. if your function is already working with “batch” with no interaction between them. Then you can compute a regular Jv and it will give the same result.
b. If there are interaction between the batches in the forward and you don’t want to see that in the Jacobian, you’ll have to do a for-loop around the Jv computation for now I’m afraid.