Hi,
the functional API works at the autograd level and so is not aware of things like batch size.
Supporting this kind of use case is planned to be supported efficiently by using vmap with the functional API.
If you’re using a regular neural net, the simplest thing to do here is just to take the diagonal of the full jacobian you get (using dim1=0 and dim2=2). It does lead to useless computations but not that much.
In particular, if you call backward on your model (in any way) all the ops will work with the full batch size of the forward.
Your proposed solution will avoid these extra computations but might not be worth the loss in forward.
You can try them both and see the result. I would be interested to know which one is faster for your usecase?