Tensor

The core array type and the autograd engine. Arithmetic operators (+, -, *, /, **, @) are overloaded and differentiable; the named methods below cover reductions, shape ops, and elementwise math.

Tensor

Tensor(data, *, requires_grad=False)

An n-dimensional array that records operations for reverse-mode autodiff.

Wraps a NumPy array. When requires_grad=True, every operation builds a node in the computation graph; calling .backward() on a scalar walks that graph in reverse and fills in .grad on the leaves.

x = Tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()
y.backward()
x.grad  # array([2., 4., 6.])

Source code in minitorch/tensor.py

def __init__(self, data, *, requires_grad=False):
    if isinstance(data, np.ndarray):
        self.data = data if data.dtype in (np.float32, np.float64) else data.astype(np.float32)
    elif isinstance(data, np.floating):
        self.data = np.array(data, dtype=data.dtype)
    else:
        self.data = np.array(data, dtype=np.float32)
    self.requires_grad = requires_grad and _grad_enabled
    self.grad = None
    self._backward = lambda: None
    self._prev = set()
    self._op = ''

sum

sum(axis=None, keepdims=False)

Sum over axis (or all elements). Differentiable.

Source code in minitorch/tensor.py

def sum(self, axis=None, keepdims=False):
    """Sum over `axis` (or all elements). Differentiable."""
    data = self.data.sum(axis=axis, keepdims=keepdims)
    out = Tensor(data, requires_grad=self.requires_grad and _grad_enabled)

    def _backward():
        if self.requires_grad:
            grad = out.grad
            if axis is not None and not keepdims:
                grad = np.expand_dims(grad, axis=axis)
            _accum_grad(self, np.broadcast_to(grad, self.data.shape))

    out._backward = _backward
    out._prev = {self}
    out._op = 'sum'
    return out

mean

mean(axis=None, keepdims=False)

Mean over axis (or all elements). Differentiable.

Source code in minitorch/tensor.py

def mean(self, axis=None, keepdims=False):
    """Mean over `axis` (or all elements). Differentiable."""
    data = self.data.mean(axis=axis, keepdims=keepdims)
    out = Tensor(data, requires_grad=self.requires_grad and _grad_enabled)

    def _backward():
        if self.requires_grad:
            if axis is None:
                count = self.data.size
            elif isinstance(axis, int):
                count = self.data.shape[axis]
            else:
                count = 1
                for a in axis:
                    count *= self.data.shape[a]
            grad = out.grad
            if axis is not None and not keepdims:
                grad = np.expand_dims(grad, axis=axis)
            _accum_grad(self, np.broadcast_to(grad, self.data.shape) / count)

    out._backward = _backward
    out._prev = {self}
    out._op = 'mean'
    return out

reshape

reshape(*shape)

Return a view with a new shape. Differentiable.

Source code in minitorch/tensor.py

def reshape(self, *shape):
    """Return a view with a new shape. Differentiable."""
    data = self.data.reshape(*shape)
    out = Tensor(data, requires_grad=self.requires_grad and _grad_enabled)

    def _backward():
        if self.requires_grad:
            _accum_grad(self, out.grad.reshape(self.data.shape))

    out._backward = _backward
    out._prev = {self}
    out._op = 'reshape'
    return out

transpose

transpose(dim0=-2, dim1=-1)

Swap two axes (last two by default). Differentiable.

Source code in minitorch/tensor.py

def transpose(self, dim0=-2, dim1=-1):
    """Swap two axes (last two by default). Differentiable."""
    axes = list(range(self.data.ndim))
    axes[dim0], axes[dim1] = axes[dim1], axes[dim0]
    data = self.data.transpose(axes)
    out = Tensor(data, requires_grad=self.requires_grad and _grad_enabled)

    def _backward():
        if self.requires_grad:
            _accum_grad(self, out.grad.transpose(axes))

    out._backward = _backward
    out._prev = {self}
    out._op = 'transpose'
    return out

exp

exp()

Elementwise e ** x. Differentiable.

Source code in minitorch/tensor.py

def exp(self):
    """Elementwise `e ** x`. Differentiable."""
    out = Tensor(np.exp(self.data), requires_grad=self.requires_grad and _grad_enabled)

    def _backward():
        if self.requires_grad:
            _accum_grad(self, out.data * out.grad)

    out._backward = _backward
    out._prev = {self}
    out._op = 'exp'
    return out

backward

backward()

Run reverse-mode autodiff from this scalar, filling .grad on every leaf.

Source code in minitorch/tensor.py

def backward(self):
    """Run reverse-mode autodiff from this scalar, filling `.grad` on every leaf."""
    assert self.data.size == 1, "backward() only works on scalar tensors - call .sum() or .mean() first"
    if self.grad is None:
        self.grad = np.ones_like(self.data)

    # iterative post-order DFS so deep graphs don't blow the recursion limit
    topo = []
    visited = set()
    stack = [(self, False)]
    while stack:
        node, processed = stack.pop()
        if processed:
            topo.append(node)
            continue
        if node in visited:
            continue
        visited.add(node)
        stack.append((node, True))
        for child in node._prev:
            if child not in visited:
                stack.append((child, False))

    for node in reversed(topo):
        node._backward()

detach

detach()

Return a tensor sharing this data but cut out of the graph.

Source code in minitorch/tensor.py

def detach(self):
    """Return a tensor sharing this data but cut out of the graph."""
    return Tensor(self.data, requires_grad=False)

clone

clone()

Return a copy of this tensor's data, keeping requires_grad.

Source code in minitorch/tensor.py

def clone(self):
    """Return a copy of this tensor's data, keeping `requires_grad`."""
    return Tensor(self.data.copy(), requires_grad=self.requires_grad)

no_grad

no_grad

Context manager to disable gradient tracking during inference.

Free functions

cat

cat(tensors, axis=0)

Concatenate tensors along an existing axis. Differentiable.

Source code in minitorch/tensor.py

def cat(tensors, axis=0):
    """Concatenate tensors along an existing axis. Differentiable."""
    data = np.concatenate([t.data for t in tensors], axis=axis)
    any_grad = any(t.requires_grad for t in tensors)
    out = Tensor(data, requires_grad=any_grad and _grad_enabled)

    def _backward():
        sizes = [t.data.shape[axis] for t in tensors]
        grads = np.split(out.grad, np.cumsum(sizes[:-1]), axis=axis)
        for t, g in zip(tensors, grads):
            if t.requires_grad:
                _accum_grad(t, g)

    out._backward = _backward
    out._prev = set(tensors)
    out._op = 'cat'
    return out

stack

stack(tensors, axis=0)

Stack tensors along a new axis. Differentiable.

Source code in minitorch/tensor.py

def stack(tensors, axis=0):
    """Stack tensors along a new axis. Differentiable."""
    data = np.stack([t.data for t in tensors], axis=axis)
    any_grad = any(t.requires_grad for t in tensors)
    out = Tensor(data, requires_grad=any_grad and _grad_enabled)

    def _backward():
        for i, t in enumerate(tensors):
            if t.requires_grad:
                slices = [slice(None)] * out.grad.ndim
                slices[axis] = i
                _accum_grad(t, out.grad[tuple(slices)])

    out._backward = _backward
    out._prev = set(tensors)
    out._op = 'stack'
    return out