[ADD] Sub-sampling in ``BatchGrad`` extension

f-dangel / backpack

BackPACK - a backpropagation package built on top of PyTorch which efficiently computes quantities other than the gradient.

https://backpack.pt/

MIT License

561 stars 55 forks source link

[ADD] Sub-sampling in ``BatchGrad`` extension #200

Closed f-dangel closed 3 years ago

f-dangel commented 3 years ago

This PR extends BackPACK's individual gradient extension and allows specifying the mini-batch samples for which the quantity will be extracted (#12), rather than evaluating individual gradients for every sample in the mini-batch.

Auxiliary:

Improve type annotations and code of abstract module extension class
Support sub-sampling in the test suite and reduce duplication
Rewrite BatchGrad tests using pytest fixtures