Custom autograd functions

Define gradients directly when standard operators are inefficient, unstable, or need to connect to custom hardware.

Define a Function

import tensorplay as tp
from tensorplay.autograd import Function

class MyExp(Function):
    @staticmethod
    def forward(ctx, value):
        result = value.exp()
        ctx.save_for_backward(result)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        result, = ctx.saved_tensors
        return grad_output * result

x = tp.randn(3, requires_grad=True)
y = MyExp.apply(x)
y.sum().backward()

When to use it

Combine gradient work for efficiency.
Implement numerically stable forms.
Connect hardware with dedicated differentiation logic.
Provide surrogate gradients for quantization and other non-differentiable operations.

The ctx object

save_for_backward stores tensors for backward.
saved_tensors retrieves them.
mark_dirty records in-place modification.
mark_non_differentiable marks outputs that do not require gradients.