Skip to content

Commit

Permalink
(doc/fluxion/adapter) add/convert docstrings to mkdocstrings format
Browse files Browse the repository at this point in the history
  • Loading branch information
Laurent2916 committed Feb 1, 2024
1 parent 82264f7 commit 58aeb6d
Showing 1 changed file with 28 additions and 0 deletions.
28 changes: 28 additions & 0 deletions src/refiners/fluxion/adapters/adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@


class Adapter(Generic[T]):
"""Base class for adapters.
An Adapter modifies the structure of a [`Module`][refiners.fluxion.layers.Module]
(typically by adding, removing or replacing layers), to adapt it to a new task.
"""

# we store _target into a one element list to avoid pytorch thinking it is a submodule
_target: "list[T]"

Expand All @@ -17,10 +23,20 @@ def __init_subclass__(cls, **kwargs: Any) -> None:

@property
def target(self) -> T:
"""The target of the adapter."""
return self._target[0]

@contextlib.contextmanager
def setup_adapter(self, target: T) -> Iterator[None]:
"""Setup the adapter.
This method should be called by the constructor of the adapter.
It sets the target of the adapter and ensures that the adapter
is not a submodule of the target.
Args:
target: The target of the adapter.
"""
assert isinstance(self, fl.Chain)
assert (not hasattr(self, "_modules")) or (
len(self) == 0
Expand All @@ -37,6 +53,13 @@ def setup_adapter(self, target: T) -> Iterator[None]:
target._can_refresh_parent = _old_can_refresh_parent

def inject(self: TAdapter, parent: fl.Chain | None = None) -> TAdapter:
"""Inject the adapter.
This method replaces the target of the adapter by the adapter inside the parent of the target.
Args:
parent: The parent to inject the adapter into, if the target doesn't have a parent.
"""
assert isinstance(self, fl.Chain)

if (parent is None) and isinstance(self.target, fl.ContextModule):
Expand All @@ -62,6 +85,11 @@ def inject(self: TAdapter, parent: fl.Chain | None = None) -> TAdapter:
return self

def eject(self) -> None:
"""Eject the adapter.
This method is the inverse of [`inject`][refiners.fluxion.adapters.Adapter.inject],
and should leave the target in the same state as before the injection.
"""
assert isinstance(self, fl.Chain)

# In general, the "actual target" is the target.
Expand Down

0 comments on commit 58aeb6d

Please sign in to comment.