Tensor & Autograd

simplegrad.core.autograd.Tensor

N-dimensional array with automatic differentiation support.

Wraps a numpy array and records operations into a dynamic computation graph. Call .backward() on a scalar output to propagate gradients back to all leaf tensors with comp_grad=True.

Parameters:

Name	Type	Description	Default
values	ndarray | list | None	Initial data. Converted to a numpy array of the given dtype.	None
comp_grad	bool	Enable gradient tracking. Defaults to the global _COMP_GRAD flag.	None
label	str | None	Optional name shown in computation graph visualizations.	None
dtype	str | None	Data type string (e.g. "float32"). Defaults to "float32".	None

Source code in simplegrad/core/autograd.py

class Tensor:
    """N-dimensional array with automatic differentiation support.

    Wraps a numpy array and records operations into a dynamic computation graph.
    Call `.backward()` on a scalar output to propagate gradients back to all
    leaf tensors with `comp_grad=True`.

    Args:
        values: Initial data. Converted to a numpy array of the given dtype.
        comp_grad: Enable gradient tracking. Defaults to the global `_COMP_GRAD` flag.
        label: Optional name shown in computation graph visualizations.
        dtype: Data type string (e.g. ``"float32"``). Defaults to ``"float32"``.
    """

    def __init__(
        self,
        values: np.ndarray | list | None = None,
        comp_grad: bool = None,
        label: str | None = None,
        dtype: str | None = None,
    ) -> None:
        self.dtype = dtype if dtype is not None else "float32"
        if values is None:
            values = np.array([])
        self.values = as_array(values, self.dtype)
        self.shape = self.values.shape
        self._forward_fn = None

        self.label = label
        self.prev = set()
        self.oper = None
        self.comp_grad = comp_grad if comp_grad is not None else get_comp_grad()
        self.is_leaf = True
        self.grad = None
        self.backward_step = lambda: None
        self.group: tuple[str, int] | None = get_current_group()

    @classmethod
    def deferred(
        cls, forward_fn: Callable[[], np.ndarray], shape: tuple, dtype: str = "float32"
    ) -> "Tensor":
        """Create an unrealized tensor that defers computation to ``.realize()``.

        Used internally by ``_create_op_result`` when lazy mode is active. The
        tensor is a shell — ``values`` is ``None`` and ``shape`` is known —
        until ``.realize()`` walks the graph and calls ``forward_fn``.

        Args:
            forward_fn: Zero-argument callable that returns a numpy ndarray.
                Captures input tensors by reference, so it stays valid until
                ``.realize()`` populates their ``.values``.
            shape: Output shape. Available immediately without executing the op.
            dtype: Data type string. Defaults to ``"float32"``.

        Returns:
            An unrealized Tensor with ``values=None``.
        """
        t = cls.__new__(cls)
        t.dtype = dtype
        t.values = None
        t.shape = shape
        t._forward_fn = forward_fn
        t.label = None
        t.prev = set()
        t.oper = None
        t.comp_grad = get_comp_grad()
        t.is_leaf = True
        t.grad = None
        t.backward_step = lambda: None
        t.group = get_current_group()
        return t

    def convert_to(self, dtype: str, inplace: bool = True) -> "Tensor" | None:
        """Convert tensor values (and gradients) to a different dtype.

        Args:
            dtype: Target dtype string (e.g. ``"float64"``).
            inplace: Modify this tensor in-place if True, else return a new tensor.

        Returns:
            New Tensor if ``inplace=False``, else None.
        """
        if self.values is None:
            raise RuntimeError(
                "Cannot convert dtype of an unrealized tensor. Call .realize() first."
            )
        if inplace:
            self.dtype = dtype
            self.values = convert_to_dtype(self.values, dtype)
            if self.grad is not None:
                self.grad = convert_to_dtype(self.grad, dtype)
        else:
            new_tensor = Tensor(
                values=convert_to_dtype(self.values, dtype),
                comp_grad=self.comp_grad,
                label=self.label,
            )
            if self.grad is not None:
                new_tensor.grad = convert_to_dtype(self.grad, dtype)
            return new_tensor

    def __len__(self) -> int:
        if self.values is None:
            raise RuntimeError("Cannot get length of an unrealized tensor. Call .realize() first.")
        return len(self.values)

    def __eq__(self, other: any) -> bool:
        return id(self) == id(other)

    def __hash__(self) -> int:
        return hash(id(self))

    def __getitem__(self, idxs) -> tuple:
        if self.values is None:
            raise RuntimeError("Cannot index an unrealized tensor. Call .realize() first.")
        return (
            self.values.__getitem__(idxs),
            self.grad.__getitem__(idxs) if self.grad is not None else None,
        )

    def __iter__(self):
        if self.values is None:
            raise RuntimeError("Cannot iterate over an unrealized tensor. Call .realize() first.")
        return self.values.__iter__()

    def __str__(self):
        if self.values is None:
            return (
                f"Tensor '{self.label}' [unrealized]\n"
                f"shape: {self.shape}\n"
                f"is_leaf: {self.is_leaf}\n"
                f"dtype: {self.dtype}\n"
                f"comp_grad: {self.comp_grad}"
            )
        if self.grad is not None:
            grad_info = f"\ngrad:\n{self.grad}"
        else:
            grad_info = "\ngrad: None"
        return (
            f"Tensor '{self.label}'\nshape: {self.shape}\nis_leaf: {self.is_leaf}\n"
            f"dtype: {self.dtype}\ncomp_grad: {self.comp_grad}\nvalues:\n{self.values}{grad_info}"
        )

    def zero_grad(self):
        """Zero gradients on all leaf tensors in the computation graph."""
        if self.comp_grad and self.is_leaf:
            self.grad = np.zeros(self.shape)
        for t in self.prev:
            t.zero_grad()

    def realize(self) -> "Tensor":
        """Execute all pending forward computations in the computation graph.

        Walks the graph in topological order (inputs before outputs) and
        executes any stored ``_forward_fn`` callables, filling in
        ``tensor.values`` for every unrealized tensor. After this call,
        ``self.values`` and the values of every upstream tensor are guaranteed
        to be non-None.

        In eager mode this is a no-op — all tensors are already realized.

        Returns:
            ``self``, so you can chain: ``loss = model(x).realize()``.
        """
        topo: list[Tensor] = []
        visited: set[int] = set()

        def _build_topo(t: Tensor) -> None:
            if id(t) in visited:
                return
            visited.add(id(t))
            for parent in t.prev:
                _build_topo(parent)
            topo.append(t)

        _build_topo(self)

        for t in topo:
            if t._forward_fn is not None:
                t.values = t._forward_fn()
                t.shape = t.values.shape
                t._forward_fn = None

        return self

    def _check_can_backward(self):
        """Raise if backward() cannot be called on this tensor."""
        if not self.comp_grad:
            raise RuntimeError(
                f"Cannot call backward() on tensor {self.label or ''} with comp_grad=False."
            )
        if self.grad is not None and not self.is_leaf:
            raise RuntimeError(
                "backward() can only be called once on non-leaf tensors, or you need to use retain_grad()"
            )
        if self.values is not None and self.values.size == 0:
            raise RuntimeError("Cannot call backward() on an empty tensor")

    def backward(self) -> None:
        """Run backpropagation from this tensor.

        If the tensor is unrealized (lazy mode), calls ``.realize()``
        automatically before running backprop. This means you can always call
        ``.backward()`` directly without a separate ``.realize()`` step.

        Computes gradients for all leaf tensors in the graph with ``comp_grad=True``.
        This tensor's gradient is initialized to ones (assumes scalar loss).
        Non-leaf gradients are freed after the backward pass.

        Raises:
            RuntimeError: If ``comp_grad=False``, tensor is empty, or backward has
                already been called on this non-leaf tensor.
        """
        self.realize()
        self._check_can_backward()

        topo_order = []
        visited = set()

        def build_topo(tensor):
            if tensor not in visited:
                visited.add(tensor)
                for t in tensor.prev:
                    build_topo(t)
                topo_order.append(tensor)

        build_topo(self)

        self.grad = np.ones(self.values.shape)
        for t in reversed(topo_order):
            t.backward_step()

        for t in topo_order:
            if not t.is_leaf and t != self:
                t.grad = None

    def _init_grad_if_needed(self) -> None:
        """Initialize gradient array to zeros if not yet set."""
        if self.grad is None:
            self.grad = np.zeros(self.shape)

    def _reduce_broadcasted_dims(self, delta: np.ndarray) -> np.ndarray:
        """Sum gradient over broadcast dimensions to match this tensor's shape."""
        while delta.ndim > self.grad.ndim:
            delta = delta.sum(axis=0)
        for i, (d, g) in enumerate(zip(delta.shape, self.grad.shape)):
            if d != g:
                delta = delta.sum(axis=i, keepdims=True)
        return delta

    def __add__(self, other: int | float | "Tensor") -> "Tensor":
        """Add a scalar or tensor element-wise."""
        if isinstance(other, (int, float)):
            return _AddScalar.apply(self, other, oper=f"+({other:.2f})")
        if not isinstance(other, Tensor):
            raise ValueError(f"Wrong operand type: {type(other)}")
        if self == other:
            return self * 2
        return _Add.apply(self, other)

    def __mul__(self, other: int | float | "Tensor") -> "Tensor":
        """Multiply by a scalar or tensor element-wise."""
        if isinstance(other, (int, float)):
            return _MulScalar.apply(self, other, oper=f"*({other:.2f})")
        if not isinstance(other, Tensor):
            raise ValueError(f"Wrong operand type: {type(other)}")
        if self == other:
            return self**2
        return _Mul.apply(self, other)

    def __pow__(self, other: int | float) -> "Tensor":
        """Raise to a scalar power element-wise."""
        if not isinstance(other, (float, int)):
            raise ValueError(f"Only 'float' or 'int' exponents are supported, got {type(other)}")
        if not is_lazy() and np.any(self.values < 0) and not float(other).is_integer():
            raise ValueError(
                f"Invalid: {self.label if self.label else 'Tensor'} ** {other} would be complex."
            )
        return _Pow.apply(self, other, oper=f"^{other:.2f}")

    def __matmul__(self, other: "Tensor") -> "Tensor":
        """Matrix multiply with another tensor."""
        if not isinstance(other, Tensor):
            raise ValueError(f"Only 'Tensor' operands are supported for matmul, got {type(other)}")
        return _Matmul.apply(self, other)

    @property
    def T(self) -> "Tensor":
        """Transpose of the tensor."""
        return _Transpose.apply(self)

    def __sub__(self, other: int | float | "Tensor") -> "Tensor":
        return self + other * -1

    def __truediv__(self, other: int | float | "Tensor") -> "Tensor":
        return self * other**-1

    def __radd__(self, other: int | float | "Tensor") -> "Tensor":
        return self + other

    def __rsub__(self, other: int | float | "Tensor") -> "Tensor":
        return other + self * -1

    def __rmul__(self, other: int | float | "Tensor") -> "Tensor":
        return self * other

    def __rtruediv__(self, other: int | float | "Tensor") -> "Tensor":
        return other * self**-1

    def __neg__(self) -> "Tensor":
        return self * -1

    @property
    def _str_id(self) -> str:
        return str(id(self))

T: 'Tensor' property

Transpose of the tensor.

__add__(other: int | float | 'Tensor') -> 'Tensor'

Add a scalar or tensor element-wise.

Source code in simplegrad/core/autograd.py

def __add__(self, other: int | float | "Tensor") -> "Tensor":
    """Add a scalar or tensor element-wise."""
    if isinstance(other, (int, float)):
        return _AddScalar.apply(self, other, oper=f"+({other:.2f})")
    if not isinstance(other, Tensor):
        raise ValueError(f"Wrong operand type: {type(other)}")
    if self == other:
        return self * 2
    return _Add.apply(self, other)

__matmul__(other: 'Tensor') -> 'Tensor'

Matrix multiply with another tensor.

Source code in simplegrad/core/autograd.py

def __matmul__(self, other: "Tensor") -> "Tensor":
    """Matrix multiply with another tensor."""
    if not isinstance(other, Tensor):
        raise ValueError(f"Only 'Tensor' operands are supported for matmul, got {type(other)}")
    return _Matmul.apply(self, other)

__mul__(other: int | float | 'Tensor') -> 'Tensor'

Multiply by a scalar or tensor element-wise.

Source code in simplegrad/core/autograd.py

def __mul__(self, other: int | float | "Tensor") -> "Tensor":
    """Multiply by a scalar or tensor element-wise."""
    if isinstance(other, (int, float)):
        return _MulScalar.apply(self, other, oper=f"*({other:.2f})")
    if not isinstance(other, Tensor):
        raise ValueError(f"Wrong operand type: {type(other)}")
    if self == other:
        return self**2
    return _Mul.apply(self, other)

__pow__(other: int | float) -> 'Tensor'

Raise to a scalar power element-wise.

Source code in simplegrad/core/autograd.py

def __pow__(self, other: int | float) -> "Tensor":
    """Raise to a scalar power element-wise."""
    if not isinstance(other, (float, int)):
        raise ValueError(f"Only 'float' or 'int' exponents are supported, got {type(other)}")
    if not is_lazy() and np.any(self.values < 0) and not float(other).is_integer():
        raise ValueError(
            f"Invalid: {self.label if self.label else 'Tensor'} ** {other} would be complex."
        )
    return _Pow.apply(self, other, oper=f"^{other:.2f}")

backward() -> None

Run backpropagation from this tensor.

If the tensor is unrealized (lazy mode), calls .realize() automatically before running backprop. This means you can always call .backward() directly without a separate .realize() step.

Computes gradients for all leaf tensors in the graph with comp_grad=True. This tensor's gradient is initialized to ones (assumes scalar loss). Non-leaf gradients are freed after the backward pass.

Raises:

Type	Description
RuntimeError	If comp_grad=False, tensor is empty, or backward has already been called on this non-leaf tensor.

Source code in simplegrad/core/autograd.py

def backward(self) -> None:
    """Run backpropagation from this tensor.

    If the tensor is unrealized (lazy mode), calls ``.realize()``
    automatically before running backprop. This means you can always call
    ``.backward()`` directly without a separate ``.realize()`` step.

    Computes gradients for all leaf tensors in the graph with ``comp_grad=True``.
    This tensor's gradient is initialized to ones (assumes scalar loss).
    Non-leaf gradients are freed after the backward pass.

    Raises:
        RuntimeError: If ``comp_grad=False``, tensor is empty, or backward has
            already been called on this non-leaf tensor.
    """
    self.realize()
    self._check_can_backward()

    topo_order = []
    visited = set()

    def build_topo(tensor):
        if tensor not in visited:
            visited.add(tensor)
            for t in tensor.prev:
                build_topo(t)
            topo_order.append(tensor)

    build_topo(self)

    self.grad = np.ones(self.values.shape)
    for t in reversed(topo_order):
        t.backward_step()

    for t in topo_order:
        if not t.is_leaf and t != self:
            t.grad = None

convert_to(dtype: str, inplace: bool = True) -> 'Tensor' | None

Convert tensor values (and gradients) to a different dtype.

Parameters:

Name	Type	Description	Default
dtype	str	Target dtype string (e.g. "float64").	required
inplace	bool	Modify this tensor in-place if True, else return a new tensor.	True

Returns:

Type	Description
'Tensor' | None	New Tensor if inplace=False, else None.

Source code in simplegrad/core/autograd.py

def convert_to(self, dtype: str, inplace: bool = True) -> "Tensor" | None:
    """Convert tensor values (and gradients) to a different dtype.

    Args:
        dtype: Target dtype string (e.g. ``"float64"``).
        inplace: Modify this tensor in-place if True, else return a new tensor.

    Returns:
        New Tensor if ``inplace=False``, else None.
    """
    if self.values is None:
        raise RuntimeError(
            "Cannot convert dtype of an unrealized tensor. Call .realize() first."
        )
    if inplace:
        self.dtype = dtype
        self.values = convert_to_dtype(self.values, dtype)
        if self.grad is not None:
            self.grad = convert_to_dtype(self.grad, dtype)
    else:
        new_tensor = Tensor(
            values=convert_to_dtype(self.values, dtype),
            comp_grad=self.comp_grad,
            label=self.label,
        )
        if self.grad is not None:
            new_tensor.grad = convert_to_dtype(self.grad, dtype)
        return new_tensor

deferred(forward_fn: Callable[[], np.ndarray], shape: tuple, dtype: str = 'float32') -> 'Tensor' classmethod

Create an unrealized tensor that defers computation to .realize().

Used internally by _create_op_result when lazy mode is active. The tensor is a shell — values is None and shape is known — until .realize() walks the graph and calls forward_fn.

Parameters:

Name	Type	Description	Default
forward_fn	Callable[[], ndarray]	Zero-argument callable that returns a numpy ndarray. Captures input tensors by reference, so it stays valid until .realize() populates their .values.	required
shape	tuple	Output shape. Available immediately without executing the op.	required
dtype	str	Data type string. Defaults to "float32".	'float32'

Returns:

Type	Description
'Tensor'	An unrealized Tensor with values=None.

Source code in simplegrad/core/autograd.py

@classmethod
def deferred(
    cls, forward_fn: Callable[[], np.ndarray], shape: tuple, dtype: str = "float32"
) -> "Tensor":
    """Create an unrealized tensor that defers computation to ``.realize()``.

    Used internally by ``_create_op_result`` when lazy mode is active. The
    tensor is a shell — ``values`` is ``None`` and ``shape`` is known —
    until ``.realize()`` walks the graph and calls ``forward_fn``.

    Args:
        forward_fn: Zero-argument callable that returns a numpy ndarray.
            Captures input tensors by reference, so it stays valid until
            ``.realize()`` populates their ``.values``.
        shape: Output shape. Available immediately without executing the op.
        dtype: Data type string. Defaults to ``"float32"``.

    Returns:
        An unrealized Tensor with ``values=None``.
    """
    t = cls.__new__(cls)
    t.dtype = dtype
    t.values = None
    t.shape = shape
    t._forward_fn = forward_fn
    t.label = None
    t.prev = set()
    t.oper = None
    t.comp_grad = get_comp_grad()
    t.is_leaf = True
    t.grad = None
    t.backward_step = lambda: None
    t.group = get_current_group()
    return t

realize() -> 'Tensor'

Execute all pending forward computations in the computation graph.

Walks the graph in topological order (inputs before outputs) and executes any stored _forward_fn callables, filling in tensor.values for every unrealized tensor. After this call, self.values and the values of every upstream tensor are guaranteed to be non-None.

In eager mode this is a no-op — all tensors are already realized.

Returns:

Type	Description
'Tensor'	self, so you can chain: loss = model(x).realize().

Source code in simplegrad/core/autograd.py

def realize(self) -> "Tensor":
    """Execute all pending forward computations in the computation graph.

    Walks the graph in topological order (inputs before outputs) and
    executes any stored ``_forward_fn`` callables, filling in
    ``tensor.values`` for every unrealized tensor. After this call,
    ``self.values`` and the values of every upstream tensor are guaranteed
    to be non-None.

    In eager mode this is a no-op — all tensors are already realized.

    Returns:
        ``self``, so you can chain: ``loss = model(x).realize()``.
    """
    topo: list[Tensor] = []
    visited: set[int] = set()

    def _build_topo(t: Tensor) -> None:
        if id(t) in visited:
            return
        visited.add(id(t))
        for parent in t.prev:
            _build_topo(parent)
        topo.append(t)

    _build_topo(self)

    for t in topo:
        if t._forward_fn is not None:
            t.values = t._forward_fn()
            t.shape = t.values.shape
            t._forward_fn = None

    return self

zero_grad()

Zero gradients on all leaf tensors in the computation graph.

Source code in simplegrad/core/autograd.py

def zero_grad(self):
    """Zero gradients on all leaf tensors in the computation graph."""
    if self.comp_grad and self.is_leaf:
        self.grad = np.zeros(self.shape)
    for t in self.prev:
        t.zero_grad()

---

Execution mode

simplegrad.core.autograd.no_grad()

Context manager that disables gradient computation globally.

Source code in simplegrad/core/autograd.py

@contextmanager
def no_grad():
    """Context manager that disables gradient computation globally."""
    global _COMP_GRAD
    prev_comp_grad = _COMP_GRAD
    _COMP_GRAD = False
    try:
        yield
    finally:
        _COMP_GRAD = prev_comp_grad

simplegrad.core.autograd.lazy()

Context manager that activates lazy execution mode for the enclosed block.

Inside the block, tensor operations record their computation without running numpy. Call .realize() on the output tensor (or call .backward()) to execute the full graph.

Example

with sg.lazy(): ... out = model(x) ... out.realize()

Source code in simplegrad/core/autograd.py

@contextmanager
def lazy():
    """Context manager that activates lazy execution mode for the enclosed block.

    Inside the block, tensor operations record their computation without
    running numpy. Call ``.realize()`` on the output tensor (or call
    ``.backward()``) to execute the full graph.

    Example:
        >>> with sg.lazy():
        ...     out = model(x)
        ... out.realize()
    """
    global _LAZY_MODE
    prev = _LAZY_MODE
    _LAZY_MODE = True
    try:
        yield
    finally:
        _LAZY_MODE = prev

simplegrad.core.autograd.set_mode(mode: str) -> None

Set the global execution mode persistently.

Prefer the lazy() context manager for scoped control. Use set_mode only when you need the mode to persist across multiple function calls or modules.

Parameters:

Name	Type	Description	Default
mode	str	Either "eager" (default, execute numpy immediately) or "lazy" (defer execution until .realize() is called).	required

Raises:

Type	Description
ValueError	If mode is not "eager" or "lazy".

Example

sg.set_mode("lazy") out = model(x) out.realize() sg.set_mode("eager")

Source code in simplegrad/core/autograd.py

def set_mode(mode: str) -> None:
    """Set the global execution mode persistently.

    Prefer the ``lazy()`` context manager for scoped control. Use
    ``set_mode`` only when you need the mode to persist across multiple
    function calls or modules.

    Args:
        mode: Either ``"eager"`` (default, execute numpy immediately) or
            ``"lazy"`` (defer execution until ``.realize()`` is called).

    Raises:
        ValueError: If ``mode`` is not ``"eager"`` or ``"lazy"``.

    Example:
        >>> sg.set_mode("lazy")
        >>> out = model(x)
        >>> out.realize()
        >>> sg.set_mode("eager")
    """
    global _LAZY_MODE
    if mode == "lazy":
        _LAZY_MODE = True
    elif mode == "eager":
        _LAZY_MODE = False
    else:
        raise ValueError(f"Unknown mode '{mode}'. Expected 'eager' or 'lazy'.")

simplegrad.core.autograd.is_lazy() -> bool

Return True if lazy execution mode is currently active.

In lazy mode, tensor operations build a deferred computation graph instead of executing numpy immediately. Actual computation is triggered by calling .realize() on a tensor or by calling .backward(), which auto-realizes before running backprop.

Returns:

Type	Description
bool	True if lazy mode is active, False for eager mode (the default).

Source code in simplegrad/core/autograd.py

def is_lazy() -> bool:
    """Return True if lazy execution mode is currently active.

    In lazy mode, tensor operations build a deferred computation graph
    instead of executing numpy immediately. Actual computation is
    triggered by calling ``.realize()`` on a tensor or by calling
    ``.backward()``, which auto-realizes before running backprop.

    Returns:
        True if lazy mode is active, False for eager mode (the default).
    """
    return _LAZY_MODE