Class _Functional

axlearn/common/module.py:1152–1200

A pure functional call to `method_fn`.

        Returns:
            The method output.

        Raises:
            ValueError: If invoking from outside of an InvocationContext, or a context with invalid
                module path.
        """
        return self.forward(*args, **kwargs)


@functools.partial(flax_struct.dataclass, frozen=False)
class _Functional:  # pytype: disable=invalid-annotation
    """A pure functional call to `method_fn`."""

    # The function to call.
    method_fn: Callable = flax_struct.field(pytree_node=False)
    # The context to call method_fn in.
    # This will be copied to prevent method_fn from mutating the original.
    context: InvocationContext = flax_struct.field(
        pytree_node=True
    )  # pytype: disable=invalid-annotation
    # Whether to require that context.parent is current_context().
    require_parent: bool = flax_struct.field(pytree_node=False)
    # Whether to copy the argument pytrees to prevent method_fn from mutating the original.
    copy_args_tree: bool = flax_struct.field(pytree_node=False, default=True)

    def __call__(self, *args, **kwargs) -> tuple[Any, OutputCollection]:
        """Invokes method_fn in a pure functional fashion.

        The invocation will not depend on external inputs or have any side effects. The results only
        depend on the given inputs. All outputs are reflected in the return value.

        Args:
            *args: Positional arguments to method_fn.
            **kwargs: Keyword arguments to method_fn.

        Returns:
            (method_outputs, output_collection), where
            - method_outputs are the return value of the method.
            - output_collection is an OutputCollection containing summaries and state updates.

        Raises:
            ValueError: If there are circular references in args, kwargs, or context.
        """
        call = getattr(self.method_fn, "__qualname__", None) or getattr(self.method_fn, "__name__")
        logging.vlog(1, "functional: %s.%s (*%s, **%s)", call, self.method_fn, args, kwargs)

        # Some badly behaved tests call F() with an InvocationContext.state that contains
        # circular references.
        # This results in a cryptic error that doesn't make the root cause obvious.
        # So we raise a clearer error explicitly.
        raise_for_cycles(dict(context=self.context, args=args, kwargs=kwargs))
        context = self.context
        if self.copy_args_tree: