Indeed, since outputs can’t depend on any other inputs across the batch dimension, all cross-batch entries in the Jacobian are zero. So you’re still doing all the useless computations but now without a for-loop.
As mentioned above, the functional API works at the autograd level and so is not aware of things like batch size. You can think about this trick as working around this limitation by turning the batch of outputs into a sum of contributions and the batch of inputs into a single big example. Along the batch dimension, only the matching outputs and inputs will contribute a non-zero result.
Explicitly, for batches of input vectors getting transformed into batches of output vectors, the trick sums the outputs (batch_size, out) into (out) so that taking the derivative with respect to the “single big example” (batch_size, in) outputs a Jacobian with shape (out, batch_size, in). Permuting then gives you the Jacobian you want.