afnio.autodiff.function

class Function:
    """
    Base class to create custom autodiff [`Function`][.].

    To create a custom [`Function`][.], subclass this class and implement the
    [`forward`][.forward] and [`backward`][.backward] static methods. Then, to use your
    custom op in the forward pass, call the class method [`apply`][.apply].
    Do not call [`forward`][.forward] directly.

    Examples:
        >>> class Func(Function):
        >>>     @staticmethod
        >>>     def forward(ctx, x: afnio.Variable):
        >>>         reverse = x.data[::-1]
        >>>         out = afnio.Variable(data=reverse, role=x.role, requires_grad=True)
        >>>         ctx.save_for_backward(x, reverse, out)
        >>>         return out
        >>>
        >>>     @staticmethod
        >>>     def backward(ctx, grad_out):
        >>>         x, reverse, out = ctx.saved_variables
        >>>         grad = f"Here is the feedback for {x.role} (reversed): {grad_out.grad}"
        >>>         role = f"Feedback to {x.role}"
        >>>         x.grad = afnio.Variable(data=grad, role=role)
        >>>         return x.grad
        >>>
        >>> a = afnio.Variable(data="This is a string", role="Input string", requires_grad=True)
        >>> c = Func.apply(a)
    """  # noqa: E501

    def __init__(self, *args, **kwargs):
        raise RuntimeError(
            f"{self.__class__} should not be instantiated. Methods on autodiff "
            "functions are all static, so you should invoke them on the class itself. "
            "Instantiating an autodiff function is not allowed."
        )

    @staticmethod
    def forward(*args: Any, **kwargs: Any) -> Any:
        """
        Define the forward of the custom autodiff Function.

        This function is to be overridden by all subclasses.
        There are two ways to define forward:

        Examples:
            Usage 1 (Combined forward and ctx):
            >>> @staticmethod
            >>> def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
            ...     pass

            - It must accept a context `ctx` as the first argument, followed by any
            number of arguments (Variables or other types).

            Usage 2 (Separate forward and ctx):
            >>> @staticmethod
            >>> def forward(*args: Any, **kwargs: Any) -> Any:
            ...     pass

            >>> @staticmethod
            >>> def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
            ...     pass

            - The forward no longer accepts a `ctx` argument.
            - Instead, you must also override the [`afnio.autodiff.Function.setup_context`][..setup_context]
            staticmethod to handle setting up the `ctx` object. `output` is the output
            of the forward, `inputs` are a Tuple of inputs to the forward.

        The context can be used to store arbitrary data that can be then retrieved
        during the backward pass. [`Variable`][afnio.Variable]s should not be stored
        directly on `ctx`. Instead, Variables should be saved either with
        `ctx.save_for_backward` if they are intended to be used in [`backward`][..backward].
        """  # noqa: E501
        raise NotImplementedError(
            "You must implement the forward function for custom autodiff.Function."
        )

    @staticmethod
    def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> Any:
        """
        There are two ways to define the forward pass of an autodiff [`Function`][..].

        Either:

        1. Override forward with the signature `forward(ctx, *args, **kwargs)`.
           `setup_context` is not overridden. Setting up the `ctx` for backward
           happens inside the `forward`.
        2. Override forward with the signature `forward(*args, **kwargs)` and
           override `setup_context`. Setting up the `ctx` for backward happens
           inside `setup_context` (as opposed to inside the `forward`)
        """
        raise NotImplementedError("setup_context is not implemented.")

    @staticmethod
    def backward(ctx: Any, *grad_outputs: Any) -> Any:
        """
        Define a formula for differentiating the operation with backward mode
        automatic differentiation.

        This function is to be overridden by all subclasses.

        It must accept a context `ctx` as the first argument, followed by
        as many outputs as the [`forward`][..forward] returned (None will be passed in
        for non variable outputs of the forward function), and it should return as many
        variables, as there were inputs to [`forward`][..forward]. Each argument is the
        gradient w.r.t the given output, and each returned value should be the gradient
        w.r.t. the corresponding input. If an input is not a
        [`Variable`][afnio.Variable] or is a [`Variable`][afnio.Variable]
        not requiring grads, you can just pass `None` as a gradient for that input.

        The context can be used to retrieve variables saved during the forward
        pass. It also has an attribute `ctx.needs_input_grad` as a tuple of booleans
        representing whether each input needs gradient. E.g., [`backward`][..backward]
        will have `ctx.needs_input_grad[0] = True` if the first input to
        [`forward`][..forward] needs gradient computed w.r.t. the output.
        """
        raise NotImplementedError(
            "You must implement the backward for your custom autodiff.Function "
            "to use it with backward mode AD (automatic differentiation)."
        )

    @classmethod
    def apply(cls, *args, **kwargs):
        """
        Applies the forward function of the custom Function class.

        This method handles cases where [`setup_context`][..setup_context] is defined
        to set up the `ctx` (context) object separately or within the
        [`forward`][..forward] method itself.
        """

        # Serialize the function and arguments
        function_name = cls.__name__

        serialized_args = [_serialize_arg(a) for a in args]
        serialized_kwargs = {k: _serialize_arg(v) for k, v in kwargs.items()}

        # Send the RPC call to the server
        try:
            # Get the singleton websocket client
            _, ws_client = get_default_clients()

            payload = {
                "function_name": function_name,
                "grad_enabled": is_grad_enabled(),
                "args": serialized_args,
                "kwargs": serialized_kwargs,
            }
            response = run_in_background_loop(ws_client.call("run_function", payload))
            if "error" in response:
                raise RuntimeError(
                    response["error"]["data"].get("exception", response["error"])
                )

            logger.debug(f"Function instantiated and shared with the server: {cls!r}")

            # Deserialize the result
            result_data = response.get("result", {}).get("data")
            if not result_data:
                raise RuntimeError(
                    f"Server did not return any data for Function.apply pass: "
                    f"payload={payload!r}, response={response!r}"
                )

            return _deserialize_fn_output(result_data)

        except Exception as e:
            logger.error(f"Failed to run function forward pass on the server: {e}")
            raise