Add module.train() to move a module back to the autodiff backend

[laggui]

Checklist

• Confirmed that cargo run-checks command has been executed.

Related Issues/PRs

From discord:

Is there a convenient way to get the model on the auto diff backend from learner.fit or do I basically have to save and reload the model if I want to do auto diff stuff after fitting?

Changes

Added the inverse function to module.valid() <> module.train(), which moves the module and all of its sub-modules to the autodiff backend.

• Added HasAutodiffModule helper trait to associate a module with its autodiff version (otherwise module.train() would require type annotation for both the autodiff backend targeted AND the module type)
• Added require_grad field to Param so module.valid() <> module.train() is stateful

Testing

• Single test for batchnorm module to illustrate usage
• Stateful require_grad tests for module and param

[codecov]

Codecov Report

❌ Patch coverage is 85.71429% with 27 lines in your changes missing coverage. Please review.
✅ Project coverage is 65.44%. Comparing base (3af2817) to head (7644504).
⚠️ Report is 98 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-core/src/module/param/primitive.rs	20.00%	12 Missing ⚠️
crates/burn-core/src/module/param/constant.rs	25.00%	9 Missing ⚠️
crates/burn-core/src/module/param/tensor.rs	83.33%	6 Missing ⚠️

❌ Your project check has failed because the head coverage (65.44%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3975      +/-   ##
==========================================
+ Coverage   64.73%   65.44%   +0.70%     
==========================================
  Files        1178     1183       +5     
  Lines      140212   139417     -795     
==========================================
+ Hits        90770    91241     +471     
+ Misses      49442    48176    -1266     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

This PR has been marked as stale because it has not been updated for over a month