Compound Ops

Compound ops are higher-level operations that are built by composing other simplegrad primitives rather than calling NumPy directly. The @compound_op decorator and graph_group context manager tag all tensors created during such a call with a shared group identifier, so the computation graph visualiser can draw a labelled cluster around them for easier inspection.

import simplegrad as sg
from simplegrad.core import compound_op

@compound_op
def my_op(x):
    return (x * x).sum()

y = my_op(sg.Tensor([1.0, 2.0, 3.0], requires_grad=True))
sg.draw(y)  # nodes for my_op are grouped in the graph

compound_op(func: Callable) -> Callable

Decorator for functions that compose multiple simplegrad operations.

Apply to any public function that builds its output by calling other simplegrad ops rather than calling numpy directly. When the function is called, all Tensors created during its execution are tagged with a shared group identifier. The computation graph renderer draws a black-border rectangle around those nodes, labelled with the function name.

Each call to the decorated function creates a new unique group, so two calls to softmax produce two separate rectangles in the graph.

Parameters:

• func (Callable) –

  The compound function to decorate. Its __name__ is used as the cluster label.

graph_group(name: str)

Context manager that tags all Tensors created inside with a named group.

Used internally by the compound_op decorator. Every Tensor created while this context is active receives a group attribute of (name, unique_id), which the graph renderer uses to draw a labelled rectangle around all nodes belonging to the same compound operation call.

Parameters:

• name (str) –

  Display name for the group, shown as a label on the cluster rectangle in the computation graph.