class

Function

Function()

source

Base class for custom differentiable operations.

Subclass to define an operation whose forward and backward passes Lucid's autograd cannot deduce automatically — for example, an operation that wraps an external library, an op with a custom backward formula for numerical stability, or a piecewise-defined function whose gradient differs from naïve autograd.

Define :pyforward and :pybackward as @staticmethod on the subclass, then invoke the op via :pyapply (NOT by calling forward directly — that would skip autograd registration).

Parameters

None

Function itself is never instantiated. Subclasses are stateless — their behaviour is defined entirely by the forward / backward static methods, and ops are invoked through the :pyapply classmethod rather than __init__.

Attributes

forwardstaticmethod

forward(ctx, *args, **kwargs) -> Tensor or tuple[Tensor, ...]. Computes the primal value; stores anything backward will need on ctx.

backwardstaticmethod

backward(ctx, *grad_outputs) -> tuple of Tensors (or None per non-tensor input). Returns the cotangents matching the positional inputs of forward.

applyclassmethod

Bind forward to the autograd graph, returning the forward output and registering backward as the gradient closure.

Notes

The :pyFunctionCtx passed to forward carries the bookkeeping autograd needs to call backward later — saved tensors plus per-input non-differentiable flags.

Each :pyFunction defines a node in the computation graph

\mathbf{y} = f(\mathbf{x}_1, \ldots, \mathbf{x}_n; \theta)

and supplies the chain-rule contribution

\frac{\partial \mathcal{L}}{\partial \mathbf{x}_i} = \sum_j \frac{\partial \mathcal{L}}{\partial y_j} \cdot \frac{\partial y_j}{\partial x_{i}}

for each input $\mathbf{x}_i$ . The backward method implements precisely this Jacobian-vector product.

Examples

Custom ReLU with an explicit backward:
>>> import lucid
>>> from lucid.autograd import Function
>>> class MyReLU(Function):
...     @staticmethod
...     def forward(ctx, x):
...         ctx.save_for_backward(x)
...         return x.clamp(min=0.0)
...     @staticmethod
...     def backward(ctx, grad_out):
...         (x,) = ctx.saved_tensors
...         return grad_out * (x > 0).float()
>>> y = MyReLU.apply(lucid.tensor([-1.0, 2.0], requires_grad=True))
>>> y.sum().backward()

Methods (3)

static

forward

→Tensor or tuple of Tensor

forward(ctx: FunctionCtx, args: Tensor = ())

source

Compute the forward result of the custom op.

Subclasses override this static method. Save anything needed during backward on ctx (via ctx.save_for_backward(...) or by setting attributes like ctx.shape = x.shape); do not capture tensors via closure.

Parameters

ctxFunctionCtx

Per-call context; will be passed unchanged to backward.

*argsTensor= ()

Positional inputs to the custom op. Non-Tensor positional arguments are allowed but receive no gradients.

Returns

Tensor or tuple of Tensor

One or more output tensors.

static

backward

→Tensor or tuple of Tensor

backward(ctx: FunctionCtx, grad_outputs: Tensor = ())

source

Compute gradients of the loss w.r.t. each input of forward.

Subclasses override this static method. The returned tuple must contain one gradient per positional argument of forward, in the same order; non-Tensor arguments receive None.

Parameters

ctxFunctionCtx

The same context that was populated during forward.

*grad_outputsTensor= ()

Upstream gradients

\partial L / \partial y_i

, one per output that forward returned.

Returns

Tensor or tuple of Tensor

Downstream gradients matching the positional inputs of forward. Use None for inputs that have no gradient.

cls

apply

→object

apply(args: object = (), kwargs: object = {})

source

Run the function — autograd-tracked when forward requests grad.

FunctionMeta rewrites this on every concrete subclass to dispatch through forward/backward; the base implementation here exists only so Function.apply is a real attribute (matters for tooling and hasattr checks). Calling it on the base class is meaningless and raises.