onnx_light.onnx_core.gradient

Reverse-mode automatic differentiation for ONNX graphs.

This module provides two functions that compute gradient FunctionProtos from ONNX graph descriptions:

• gradient_of_nodes() – takes a list of NodeProto and metadata (inputs, initializers, xs, y, zs) and returns a FunctionProto encoding the gradient computation.

• gradient_of_function() – takes an existing FunctionProto together with xs, y, zs and returns the gradient FunctionProto.

The returned FunctionProto has:

• inputs: xs values, then zs values, then "dy" (the incoming gradient of y; pass ones_like(y) for a scalar loss).

• outputs: one gradient tensor per element of xs, named "grad_<x>".

Supported forward operators

Conv, MatMul, Gemm, Add, Sub, Mul, Div, Neg, Identity, Relu, Sigmoid, Tanh, ReduceSum, ReduceMean, Reshape, Transpose.

Example: linear regression gradient

from onnx_light.onnx_proto._helper import make_node
from onnx_light.onnx_core.gradient import gradient_of_nodes

nodes = [
    make_node("MatMul", ["X", "W"], ["mm"]),
    make_node("Add", ["mm", "b"], ["y"]),
]

grad_fn = gradient_of_nodes(
    nodes=nodes,
    inputs=["X", "W", "b"],
    initializers=[],
    xs=["W", "b"],
    y="y",
    zs=["X"],
)
# grad_fn.input  = ["W", "b", "X", "dy"]
# grad_fn.output = ["grad_W", "grad_b"]

class onnx_light.onnx_core.gradient.GradRegistry(*args, **kwargs)

Maps (domain, op_type) pairs to backward gradient functions.

default = <nanobind.nb_func object>

op_types(self, domain: str = '') → list[str]

Returns a sorted list of op_type names registered for domain.

Parameters:

domain (str) – Operator domain (default: "" for standard ONNX ops).

Returns:

Sorted op_type names registered for domain.

Return type:

list[str]