Method get_state_updates

axlearn/common/module.py:529–530

(self)

        """
        # E.g., optimizer state from a PartitionedGradientTransformation may be a tuple, so we
        # have to assign it via the parent.
        parent = self.parent
        if parent is not None:
            parent.output_collection.state_updates[self.name] = value
        self.output_collection = self.output_collection._replace(state_updates=value)