TorchJD
refactor(autogram): Remove vmap rule generation
Not sure what I'm doing here, but as far as I understand, the generated vmap rule for JacobianAccumulator is never used. I think (correct me if I'm wrong) that it would only be used if we vmapped the forward pass of a hooked model. And I don't think we will ever cover that case (it's just too much work for nothing it seems). So this is 1 less line of code, one less non-tested feature, and maybe even a lighter JacobianAccumulator node.
Codecov Report
:white_check_mark: All modified and coverable lines are covered by tests.
| Files with missing lines | Coverage Δ | |
|---|---|---|
| src/torchjd/autogram/_module_hook_manager.py | 100.00% <ø> (+100.00%) |
:arrow_up: |
... and 51 files with indirect coverage changes
:rocket: New features to boost your workflow:
- :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
![codecov[bot] avatar](https://avatars.githubusercontent.com/in/254?v=4 "Published by a pub.dev verified publisher"){kind=small}
I added this when we where having 2 engines on the same modules. This is because the forward was then run using vmap with the functional api. This is exactly what we lose with this change (but I don't have anything against it).
Ok, then I'm not sure this PR is such a good idea. I'm gonna keep this open for now.