Method set_state_update

axlearn/common/module.py:513–524

Sets the state update field of the output collection. Useful when writing a "transparent module" that needs its state to have a specific structure (e.g., a tuple) for backwards compatibility.

(self, value: Any)

            name: The output name.
            value: The output tensor.

        Raises:
            OutputConflictError: If outputs of method already exist.
        """
        if name in self.output_collection.module_outputs:
            raise OutputConflictError(f"Outputs of name '{name}' already exist")

        self.output_collection.module_outputs[name] = value

    def set_state_update(self, value: Any):
        """Sets the state update field of the output collection.

        Useful when writing a "transparent module" that needs its state to have a specific
        structure (e.g., a tuple) for backwards compatibility.
        """