leap_c.autograd.function

Classes

DiffFunction

Abstract base class for differentiable functions.

Module Contents

class leap_c.autograd.function.DiffFunction

Bases: abc.ABC

Abstract base class for differentiable functions.

Subclasses must implement forward and backward methods.

• The forward method computes outputs from inputs and returns a tuple containing a context object followed by one or more output arrays.

• The backward method receives the context and the gradients of the outputs, and returns a tuple of gradients with respect to the inputs.

A minimal subclass looks like this:

class MyFunction:
    def forward(
        self, *inputs: np.ndarray, ctx: Optional[dict] = None
    ):
        if ctx is None:
            ctx = {}
        # Store intermediate values for backward
        ctx["saved"] = inputs[0]

        # Example computation
        out1 = inputs[0] + inputs[1]
        out2 = inputs[0] * 2

        return ctx, out1, out2

    def backward(
        self, ctx: dict, *grad_outputs: np.ndarray  # type: ignore
    ):
        # Retrieve saved values
        x = ctx["saved"]

        # Example gradient computation
        grad_out1, grad_out2 = grad_outputs
        grad_x = grad_out1 + 2 * grad_out2
        grad_y = grad_out1

        return grad_x, grad_y

abstractmethod backward(ctx, *output_grads: numpy.ndarray | None)

Computes the gradient of the function with respect to its inputs.

Parameters:

• ctx – The context object returned from the forward pass.

• output_grads – Gradients with respect to each output.

Returns:

A tuple of gradients with respect to each input.

abstractmethod forward(ctx=None, *inputs: numpy.ndarray | None)

Computes the output of the function given inputs.

Parameters:

• inputs – Input arrays to the function.

• ctx – Optional context object that can be used to store intermediate values for the backward pass.

Returns:

A tuple where the first element is a context object, followed by one or more output arrays.