Class InvocationContext

axlearn/common/module.py:367–551

The invocation context for `Module.__call__()`. Attributes: name: The context name. Must be unique among sibling contexts. parent: The parent context, or None if `self` is the root context. module: The Module associated with the context. state: The state of the m

@typing.runtime_checkable  # Needed for isinstance checks to work.
class Summable(Protocol):
    # Objects of the same type which adhere to this protocol may be added.
    def __add__(self, other: T) -> T: ...


# TODO(markblee): Link to docs on invocation contexts.
@functools.partial(flax_struct.dataclass, frozen=False)
# pylint: disable-next=too-many-instance-attributes
class InvocationContext:  # pytype: disable=invalid-annotation
    """The invocation context for `Module.__call__()`.

    Attributes:
        name: The context name. Must be unique among sibling contexts.
        parent: The parent context, or None if `self` is the root context.
        module: The Module associated with the context.
        state: The state of the module.
        is_training: Whether the invocation should run in the training mode.
        prng_key: The pseudo-random number generator key (can be None if the computation does not
            require random numbers).
        output_collection: See `OutputCollection`.
    """

    name: str = flax_struct.field(pytree_node=False)
    parent: Optional["InvocationContext"] = flax_struct.field(pytree_node=True)
    module: Optional["Module"] = flax_struct.field(pytree_node=False)
    state: NestedTensor = flax_struct.field(pytree_node=True)
    is_training: bool = flax_struct.field(pytree_node=False)
    prng_key: Optional[Tensor] = flax_struct.field(pytree_node=True)
    output_collection: OutputCollection = flax_struct.field(pytree_node=True)

    def path(self):
        if self.parent is None:
            return self.name
        return self.parent.path() + "." + self.name

    # pylint: disable-next=too-many-branches
    def add_child(self, name: str, **override_kwargs) -> "InvocationContext":
        """Creates a child context with the given `name`.

        Args:
            name: The child context name. Must be unique among the siblings.
            override_kwargs: Overrides of the child context fields.

        Returns:
            The child context. Default values for fields not specified in `override_kwargs`:
            - `module` defaults to self.module.children[name];
            - `state` defaults to the state corresponding to child_context.module, that is:
              self.state[child_module.name] if child_context.module is a a child of self.module, or
              self.state if child_context.module is self.module;
            - `is_training` defaults to self.is_training;
            - `prng_key` defaults to fold_in(self.prng_key, hash(name));
            - `output_collection` defaults to self.output_collection.add_child(name).

        Raises:
            ValueError: if "parent" is specified in `override_kwargs`.
            NotImplementedError: if a field doesn't have a default value.
        """