refactor(aggregation): Add Stateful mixin

[ValerianRey]

No description provided.

[ValerianRey]

/opencode:Plan what do you think about this new structure? What would be some alternatives? Please also review the PR.

[opencode-agent]

PR Review: Add ResettableMixin

Summary

The PR introduces a ResettableMixin abstract base class that formalizes the reset() pattern for stateful aggregators. This is a clean refactor that improves code organization and type safety.

Changes Overview

File	Change
_mixins.py	NEW - Adds ResettableMixin ABC with abstract reset() method
_gradvac.py	GradVac and GradVacWeighting now inherit from ResettableMixin
_nash_mtl.py	NashMTL and _NashMTLWeighting now inherit from ResettableMixin

What's Good

1. Clean abstraction - The mixin is minimal and focused, following SRP
2. Formalizes existing pattern - Both aggregators already had reset() methods; this just makes the contract explicit
3. Type safety - Enables isinstance(obj, ResettableMixin) checks for library users
4. Forces implementation - Using @abstractmethod ensures subclasses implement reset()
5. Documentation - Clear docstring explains the purpose

Potential Alternatives

1. Protocol (Structural Subtyping)

   from typing import Protocol, runtime_checkable
   
   @runtime_checkable
   class Resettable(Protocol):
       def reset(self) -> None: ...

   • Pros: No multiple inheritance needed, duck typing friendly
   • Cons: Less explicit, runtime checkable protocols have overhead

2. Duck Typing

   if hasattr(obj, 'reset'):
       obj.reset()

   • Pros: Simple, no new classes needed
   • Cons: No IDE/type checker support, error-prone

3. Decorator Pattern

   @resettable
   class MyAggregator: ...

   • Pros: Could add behavior like auto-tracking
   • Cons: Overkill for just one method

4. Interface Segregation (More Specific)

   class EMAStateMixin(ABC):  # For GradVac's EMA state
   class OptimizationStateMixin(ABC):  # For NashMTL's solver state

   • Pros: More semantic, allows different state types
   • Cons: Might be over-engineering at this point

Minor Suggestions

1. Consider exporting the mixin - If users might want to write their own resettable aggregators:

   # In __init__.py
   from ._mixins import ResettableMixin
   __all__ += ["ResettableMixin"]

2. Add a changelog entry:

   ### Added
   - Added `ResettableMixin` to provide a formal interface for stateful aggregators.

Verdict

Approve ✅ - The implementation is clean, minimal, and follows existing patterns in the codebase. The ABC approach is appropriate here since:

• The project uses ABCs elsewhere (e.g., Aggregator, Weighting)
• It provides runtime type checking capability
• It enforces the contract through @abstractmethod

The mixin could be made public if there's user demand for creating custom resettable aggregators, but keeping it private (_mixins.py) for now is a reasonable conservative approach.

opencode session  |  github run

[ValerianRey]

Main question is: should we make this class public and add a documentation entry for it? If it's private, it's just for us, to keep code organized. If it's public, it's also a tool that users could use, but then it bloats the documentation a bit.

[PierreQuinton]

LGTM. I would make it public though: users can have a "clean" isinstance condition in their training loop if they want to check if it is resetable, instead of having to check if there is a reset method (or worse, have a list of resetable aggregators).

[ValerianRey]

Ok I made it public with a bunch of other changes. I think you're gonna like it.

[ValerianRey]

Clicking on Stateful leads to:

[PierreQuinton]

Fantastic, can merge when ready