Skip to content

Compound Ops

simplegrad.core.compound_ops.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:

Name Type Description Default
func Callable

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

 required

Returns:

Type Description
Callable

Wrapped function with identical signature and docstring.
Example

@compound_op ... def softmax(x: Tensor, dim: int | None = None) -> Tensor: ... exps = exp(x) ... return exps / sum(exps, dim)

Source code in simplegrad/core/compound_ops.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def 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.

    Args:
        func: The compound function to decorate. Its ``__name__`` is used as the cluster label.

    Returns:
        Wrapped function with identical signature and docstring.

    Example:
        >>> @compound_op
        ... def softmax(x: Tensor, dim: int | None = None) -> Tensor:
        ...     exps = exp(x)
        ...     return exps / sum(exps, dim)
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        with graph_group(func.__name__):
            return func(*args, **kwargs)

    return wrapper

simplegrad.core.compound_ops.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 Type Description Default
name str

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

 required
Example

with graph_group("softmax"): ... exps = exp(x) ... result = exps / sum(exps, dim)

Source code in simplegrad/core/compound_ops.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
@contextmanager
def 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.

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

    Example:
        >>> with graph_group("softmax"):
        ...     exps = exp(x)
        ...     result = exps / sum(exps, dim)
    """
    global _CURRENT_GROUP, _GROUP_COUNTER
    _GROUP_COUNTER += 1
    _CURRENT_GROUP = (name, _GROUP_COUNTER)
    try:
        yield
    finally:
        _CURRENT_GROUP = None