PyTorch Custom Operator API (torch/_custom_ops.py)

Registers a PyTorch custom operator and its tensor, autograd, abstract, and backward implementations. Reuse when extending PyTorch with ops that must join the dispatcher and differentiation.

# Source: https://github.com/pytorch/pytorch | License: BSD-3-Clause | Copyright 2016-     Facebook, Inc            (Adam Paszke)
# Provided via justclaw.it.com; keep this notice with the code.
# mypy: allow-untyped-defs
import inspect

from torch._custom_op.impl import (
    _custom_op_with_schema,
    _find_custom_op,
    infer_schema,
    parse_qualname,
    validate_namespace,
)
from torch.library import get_ctx


__all__ = [
    "custom_op",
    "impl",
    "impl_abstract",
    "get_ctx",
    "impl_save_for_backward",
    "impl_backward",
]


def custom_op(qualname, func_or_schema=None):
    r"""Register a new custom operator

    In PyTorch, defining an op (short for "operator") is a two step-process:
    - we need to define the op (by providing an operator name and schema)
    - we 

... (truncated -- full source via MCP)