fix: improve ForwardDiff tagging for HVP

[gdalle]

Versions

• Bump DI to v0.6.18

DI extensions

ForwardDiff:

• Get rid of the custom HVP tag for ForwardDiff-over-something
• Use a variant of shuffled_gradient instead, which makes the tag type easy to predict for preparation.
• Now the inner preparation relies on additional arguments instead of a closure, which means ReverseDiff taping won't work: I have opened Fix performance for ForwardDiff-over-ReverseDiff #603 to keep track until I add a dedicated workaround.

DI source

• Implement a new, private type of Context called PrepContext to pass preparation results as part of second-order operators.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.62%. Comparing base (2fee793) to head (ec21e89).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #596      +/-   ##
==========================================
- Coverage   97.63%   97.62%   -0.02%     
==========================================
  Files         110      110              
  Lines        5377     5382       +5     
==========================================
+ Hits         5250     5254       +4     
- Misses        127      128       +1     

Flag	Coverage Δ
DI	98.71% <100.00%> (-0.03%)	⬇️
DIT	95.57% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

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