forward_override

fastforward.forward_override.T = TypeVar('T') module-attribute

fastforward.forward_override.OverrideFn

Bases: Protocol[T]

Protocol for override functions.

Override functions can be registered to override a quantizer. These acts as dynamic decorators that can be applied an removed at runtime.

__call__(__context, __overridden_fn, __args, __kwargs)

Generic call signature.

Parameters:

Name	Type	Description	Default
__context	Any	The quantizer that is overridden by this override function	required
__overridden_fn	Callable[..., T]	The function that is overridden, this may the forward method of a quantizer, or a reference to another override. This should be called instead of the forward method on __context in order for an override to cooperate with other overrides.	required
__args	tuple[Any, ...]	The arguments passed to the quantizer.	required
__kwargs	dict[str, Any]	The keyword arguments passed to the quantizer.	required

Source code in fastforward/forward_override.py

def __call__(
    self,
    __context: Any,
    __overridden_fn: Callable[..., T],
    __args: tuple[Any, ...],
    __kwargs: dict[str, Any],
) -> T:
    """Generic call signature.

    Args:
        __context: The quantizer that is overridden by this override function
        __overridden_fn: The function that is overridden, this may the
            forward method of a quantizer, or a reference to another override.
            This should be called instead of the forward method on __context
            in order for an override to cooperate with other overrides.
        __args: The arguments passed to the quantizer.
        __kwargs: The keyword arguments passed to the quantizer.
    """
    raise NotImplementedError

fastforward.forward_override.OverrideHandle(override_map)

Handle object which that can be used to remove a function override.

Parameters:

Name	Type	Description	Default
override_map	MutableMapping[int, OverrideFn[Any]]	Mapping that stores overrides indexed by override id.	required

Source code in fastforward/forward_override.py

def __init__(self, override_map: MutableMapping[int, OverrideFn[Any]]) -> None:
    self.override_map = weakref.ref(override_map)
    self.handle_id = OverrideHandle.global_handles
    OverrideHandle.global_handles += 1

global_handles = 0 class-attribute instance-attribute

handle_id = OverrideHandle.global_handles instance-attribute

override_map = weakref.ref(override_map) instance-attribute

__enter__()

Source code in fastforward/forward_override.py

def __enter__(self) -> Self:
    return self

__exit__(type, value, traceback)

Source code in fastforward/forward_override.py

def __exit__(self, type, value, traceback) -> None:  # type: ignore[no-untyped-def]
    self.remove()

remove()

Remove override associated with this handle.

Returns:

Type	Description
OverrideFn[Any] | None	(Callable) the override function associated with this handle

Source code in fastforward/forward_override.py

def remove(self) -> OverrideFn[Any] | None:
    """Remove override associated with this handle.

    Returns:
        (Callable) the override function associated with this handle
    """
    override_map = self.override_map()
    if override_map is not None and self.handle_id in override_map:
        return override_map.pop(self.handle_id)
    return None

fastforward.forward_override.apply_overrides(context, overridden_fn, override_map)

Apply overrides in override_map to overridden_fn.

This returns a callable that, when applicable, calls functions overrides. These act similarly to decorators.

The overrides are applied in descending order of the key in override_map

Parameters:

Name	Type	Description	Default
context	Any	An object related to overridden fn that may store contextual data that is passed into the function overrides.	required
overridden_fn	Callable[..., T]	The function to override.	required
override_map	Mapping[int, OverrideFn[T]]	A mapping from int to overrides. The overrides are applied in order of the key.	required

Returns:

Type	Description
_WrappedOverriddenFn[T] | Callable[..., T]	(Callable[..., T]) wrapped overridden_fn, or overridden_fn if override_map
_WrappedOverriddenFn[T] | Callable[..., T]	is empty.

Source code in fastforward/forward_override.py

def apply_overrides(
    context: Any,
    overridden_fn: Callable[..., T],
    override_map: Mapping[int, OverrideFn[T]],
) -> _WrappedOverriddenFn[T] | Callable[..., T]:
    """Apply overrides in `override_map` to `overridden_fn`.

    This returns a callable that, when applicable, calls functions overrides.
    These act similarly to decorators.

    The overrides are applied in descending order of the key in `override_map`

    Args:
        context: An object related to overridden fn that may store
            contextual data that is passed into the function overrides.
        overridden_fn: The function to override.
        override_map: A mapping from int to overrides. The overrides are applied
            in order of the key.

    Returns:
        (Callable[..., T]) wrapped overridden_fn, or overridden_fn if override_map
        is empty.
    """
    if not override_map:
        return overridden_fn
    wrapped_fn = _WrappedOverriddenFn(
        context=context,
        overridden_fn=overridden_fn,
        override_map=override_map,
    )
    return wrapped_fn