fix(utils): fix overloaded annotations of extract_state_dict (#162)

[StefanoWoerner]

Description

Added detach_buffers keyword arg to the nn.Module version of extract_state_dict and removed with_buffers from the MetaOptimizer version of extract_state_dict.

Motivation and Context

Fixes #161

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

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

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of example)

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)

Comment: there currently are no tests for type checking, so there weren't any tests that needed updating.

[XuehaiPan]

Hi @StefanoWoerner, thanks for the fix. Thanks for pointing out the missing argument in overloads. If I understand the original purpose of writing @overload annotations correctly, the differences between the two overloaded annotations are: different in the input of target and different in the function return type.

@overload
def extract_state_dict(
    target: nn.Module,
    *,
    ...
) -> ModuleState:
    ...


@overload
def extract_state_dict(
    target: MetaOptimizer,
    *,
    ...
) -> tuple[OptState, ...]:
    ...

If you could only add detach_buffers: bool = False, and keep with_buffers: bool = True, in both overloaded annotations, we can get this merged. Thanks!

[StefanoWoerner]

Since with_buffer and detach_buffers are not used for MetaOptimizers, it doesn't make sense to specify them when calling extract_state_dict with target: MetaOptimizer. Therefore I think that leaving them out in the overload definition makes the most sense. Overload definitions with different numbers/names of arguments are actually allowed, as shown in the map example in the type hinting PEP: https://peps.python.org/pep-0484/#function-method-overloading

However, I can of course also add with_buffers and detach_buffers back in if you'd still prefer it.

[codecov]

Codecov Report

Patch and project coverage have no change.

Comparison is base (47f0109) 93.37% compared to head (ad935e6) 93.37%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #162   +/-   ##
=======================================
  Coverage   93.37%   93.37%           
=======================================
  Files          71       71           
  Lines        2684     2684           
=======================================
  Hits         2506     2506           
  Misses        178      178           

Flag	Coverage Δ
unittests	93.37% <ø> (ø)

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

Impacted Files	Coverage Δ
torchopt/utils.py	86.77% <ø> (ø)

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

[XuehaiPan]

Since with_buffer and detach_buffers are not used for MetaOptimizers, it doesn't make sense to specify them when calling extract_state_dict with target: MetaOptimizer.

Thanks for the clarification and thanks for your contribution!