[ADD] AutomaticDerivatives for autograd-based support of new layers#336
[ADD] AutomaticDerivatives for autograd-based support of new layers#336f-dangel wants to merge 12 commits into
AutomaticDerivatives for autograd-based support of new layers#336Conversation
|
@fKunstner could you have quick look, or give feedback whether you think this feature makes sense? |
There was a problem hiding this comment.
Pull request overview
This PR introduces experimental automatic derivative support for BackPACK through a new AutomaticDerivatives class that leverages torch.autograd.grad to compute derivatives automatically. This enables supporting new layers without manual derivative implementation, trading performance for convenience. The implementation currently supports VJPs (Vector-Jacobian Products) w.r.t. layer inputs and parameters, covering the most common first- and second-order extensions.
Key changes:
- New
AutomaticDerivativesclass provides autograd-based derivative computation for arbitrary layers - Comprehensive test suite demonstrates usage for GGN diagonal and batch gradient computations
- Updated type hints for better type safety in subsampling utilities
- Documentation tracking and badge updates in README
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
backpack/core/derivatives/automatic.py |
Core implementation of AutomaticDerivatives using torch.func (vjp, vmap) for automatic derivative computation |
test/test_automatic_support.py |
Test suite validating automatic derivatives against manual implementations for DiagGGN and BatchGrad extensions |
test/automatic_extensions.py |
Example extension classes demonstrating how to use AutomaticDerivatives with existing BackPACK extension framework |
test/utils/__init__.py |
Added popattr helper function for removing and returning object attributes in tests |
backpack/utils/subsampling.py |
Updated type hints to use Optional[List[int]] for better type safety |
fully_documented.txt |
Updated documentation tracking to include new files (contains reference to non-existent file) |
README.md |
Updated badges (Travis to RTD) and minimum Python version (3.8+ to 3.9+) |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This PR adds an experimental
AutomaticDerivativesclass, which implements BackPACK's core functionality usingtorch.autograd.gradgiven a module's forward function. This simplifies supporting new layers at the cost of reduced performance. Currently,AutomaticDerivativesonly implements VJPs w.r.t. layer inputs and parameters. This is enough to support the most common first- and second-order extensions.A future PR will add an example to demonstrate how to use
AutomaticDerivativesto add support for a new layer to a specific extension.