@Yaroslav_Bulatov we’re working on prototyping an auto-vectorization operation like JAX’s vmap (https://github.com/pytorch/pytorch/issues/42368 30) in the hopes that it may be able to compute per-sample-gradients (among other use cases). There’s a slight problem where it doesn’t quite work for the per-sample-gradient use case, because of how PyTorch internals are organized, but we’re figuring that out.
I’m curious if you’ve come across other demands for per-example gradients. I’ve spoken with some folks in the differential privacy space but haven’t been able to find other use cases.