Class OutputCollection

axlearn/common/module.py:251–279

Implicit outputs from module invocations. Usually users do not interact with `OutputCollection` directly, but via `InvocationContext.add_{summary, state_update, module_output}()`.

class OutputCollection(NamedTuple):
    """Implicit outputs from module invocations.

    Usually users do not interact with `OutputCollection` directly, but via
    `InvocationContext.add_{summary, state_update, module_output}()`.
    """

    summaries: NestedTensor
    state_updates: NestedTensor
    module_outputs: NestedTensor

    def __contains__(self, name: str) -> bool:
        return name in self.summaries

    def add_child(self, name: str) -> "OutputCollection":
        if not re.fullmatch("^[a-z][a-z0-9_]*$", name):
            raise ValueError(f'Invalid child name "{name}"')
        if name in self:
            raise OutputConflictError(f"{name} already present")
        child = new_output_collection()
        self.summaries[name] = child.summaries
        self.state_updates[name] = child.state_updates
        self.module_outputs[name] = child.module_outputs
        return child

    def update(self, collection: "OutputCollection"):
        self.summaries.update(**collection.summaries)
        self.state_updates.update(**collection.state_updates)
        self.module_outputs.update(**collection.module_outputs)


def new_output_collection():