freenode
AI & ML

PyTorch feature request seeks decorator API for custom autograd ops

A proposed torch.autograd.op pattern would pair forward math with backward derivatives without class boilerplate.

A feature request filed against PyTorch proposes a decorator-based functional API for custom autograd operators, aiming to replace much of the class-heavy torch.autograd.Function workflow used today.

Writing a custom differentiable op currently means subclassing Function, implementing forward and backward methods, and manually stashing tensors on a context object for the backward pass. The proposal, labeled torch.autograd.op, would let authors register a plain forward function and a matching backward via decorators, with needed forward inputs forwarded automatically instead of explicit save and restore calls.

The pitch targets friction on simple mathematical ops and claims a cleaner fit with Dynamo and AOTAutograd graph capture, while keeping dispatcher boundaries intact. The request is framed as an in-tree change, with tutorial and marketing coverage called out as part of a possible rollout. No release version or design document is attached yet; the idea remains a proposal pending design and acceptance from the autograd maintainers.