feat: make root_vjp public

[Benjamin-eecs]

Description

make root_vjp public to users

Types of changes

What types of changes does your code introduce? Put an x in all the boxes that apply:

• New feature (non-breaking change which adds core functionality)
• Documentation (update in the documentation)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide. (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly. (required for a bug fix or a new feature)
• I have updated the documentation accordingly.
• I have reformatted the code using make format. (required)
• I have checked the code using make lint. (required)
• I have ensured make test pass. (required)

[codecov]

Codecov Report

Patch coverage: 92.31% and project coverage change: -0.01 ⚠️

Comparison is base (a6a2ed1) 93.37% compared to head (2fd5166) 93.36%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #157      +/-   ##
==========================================
- Coverage   93.37%   93.36%   -0.01%     
==========================================
  Files          71       72       +1     
  Lines        2684     2696      +12     
==========================================
+ Hits         2506     2517      +11     
- Misses        178      179       +1     

Flag	Coverage Δ
unittests	93.36% <92.31%> (-0.01%)	⬇️

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

Impacted Files	Coverage Δ
torchopt/diff/implicit/decorator.py	95.73% <ø> (ø)
torchopt/diff/implicit/utils.py	90.91% <90.91%> (ø)
torchopt/diff/implicit/__init__.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.

[XuehaiPan]

Need more design for the vjp API.

• root_vjp should return a function rather than the product result.
• How do you ensure the solution argument is the root of solve_fn? Seems that the user can pass arbitrary tensors to your root_vjp API.
• Need more design for argnums and argument binding. Need more investigation for functools.partial support.