Layers

Every layer subclasses Module. Activations are also available as plain functions under minitorch.functional (aliased as F).

Linear

Linear

Linear(in_features, out_features, bias=True, init='kaiming')

Bases: Module

Fully-connected layer y = x @ W + b.

Weights use Kaiming init by default (init='xavier' for glorot). Accepts any input shape ending in in_features.

Source code in minitorch/layers.py

def __init__(self, in_features, out_features, bias=True, init='kaiming'):
    super().__init__()
    assert in_features > 0 and out_features > 0, "features must be positive"
    if init == 'xavier':
        # glorot uniform - good for sigmoid/tanh
        scale = np.sqrt(2.0 / (in_features + out_features))
    else:
        # kaiming - good for relu
        scale = np.sqrt(2.0 / in_features)
    self.weight = Tensor(
        (np.random.randn(in_features, out_features) * scale).astype(np.float32),
        requires_grad=True
    )
    self.bias = Tensor(
        np.zeros(out_features, dtype=np.float32),
        requires_grad=True
    ) if bias else None

Activations

ReLU

ReLU()

Bases: Module

Rectified linear unit, max(0, x).

Source code in minitorch/module.py

def __init__(self):
    self._training = True

GELU

GELU()

Bases: Module

Source code in minitorch/module.py

def __init__(self):
    self._training = True

Sigmoid

Sigmoid()

Bases: Module

Logistic sigmoid, 1 / (1 + e**-x).

Source code in minitorch/module.py

def __init__(self):
    self._training = True

Tanh

Tanh()

Bases: Module

Hyperbolic tangent.

Source code in minitorch/module.py

def __init__(self):
    self._training = True

Softmax

Softmax(axis=-1)

Bases: Module

Softmax over axis, turning logits into a probability distribution.

Source code in minitorch/layers.py

def __init__(self, axis=-1):
    super().__init__()
    self.axis = axis

Normalization and regularization

BatchNorm1d

BatchNorm1d(num_features, eps=1e-05, momentum=0.1)

Bases: Module

Normalize a (batch, features) input per feature, with learnable scale/shift.

Tracks running mean and variance for use at eval time.

Source code in minitorch/layers.py

def __init__(self, num_features, eps=1e-5, momentum=0.1):
    super().__init__()
    assert num_features > 0, "num_features must be positive"
    self.gamma = Tensor(np.ones(num_features, dtype=np.float32), requires_grad=True)
    self.beta = Tensor(np.zeros(num_features, dtype=np.float32), requires_grad=True)
    self.eps = eps
    self.momentum = momentum
    self.running_mean = np.zeros(num_features, dtype=np.float32)
    self.running_var = np.ones(num_features, dtype=np.float32)

LayerNorm

LayerNorm(num_features, eps=1e-05)

Bases: Module

Normalize over the last dimension. Works on any (..., num_features) input.

Source code in minitorch/layers.py

def __init__(self, num_features, eps=1e-5):
    super().__init__()
    assert num_features > 0, "num_features must be positive"
    self.gamma = Tensor(np.ones(num_features, dtype=np.float32), requires_grad=True)
    self.beta = Tensor(np.zeros(num_features, dtype=np.float32), requires_grad=True)
    self.eps = eps

Dropout

Dropout(p=0.5)

Bases: Module

Zero each activation with probability p during training, scaling the rest.

A no-op in eval mode (module.eval()).

Source code in minitorch/layers.py

def __init__(self, p=0.5):
    super().__init__()
    assert 0.0 <= p < 1.0, "dropout probability must be in [0, 1)"
    self.p = p

Embedding

Embedding

Embedding(num_embeddings, embed_dim)

Bases: Module

Lookup table mapping integer ids to dense vectors.

Source code in minitorch/layers.py

def __init__(self, num_embeddings, embed_dim):
    super().__init__()
    assert num_embeddings > 0 and embed_dim > 0, "sizes must be positive"
    self.weight = Tensor(
        (np.random.randn(num_embeddings, embed_dim) * 0.02).astype(np.float32),
        requires_grad=True
    )

Convolution

Conv2d

Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0)

Bases: Module

2D convolution over (N, C, H, W) input, implemented with im2col.

Source code in minitorch/conv.py

def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
    super().__init__()
    assert in_channels > 0 and out_channels > 0, "channels must be positive"
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
    self.stride = stride
    self.padding = padding

    kh, kw = self.kernel_size
    assert kh > 0 and kw > 0, "kernel_size must be positive"
    fan_in = in_channels * kh * kw
    scale = np.sqrt(2.0 / fan_in)
    self.weight = Tensor(
        (np.random.randn(out_channels, in_channels * kh * kw) * scale).astype(np.float32),
        requires_grad=True
    )
    self.bias = Tensor(np.zeros(out_channels, dtype=np.float32), requires_grad=True)

MaxPool2d

MaxPool2d(kernel_size, stride=None)

Bases: Module

Max pooling over (N, C, H, W) input using strided windows.

Source code in minitorch/conv.py

def __init__(self, kernel_size, stride=None):
    super().__init__()
    self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
    self.stride = stride if stride is not None else self.kernel_size[0]

Flatten

Flatten()

Bases: Module

Flatten all dims except the batch dim into one.

Source code in minitorch/conv.py

def __init__(self):
    super().__init__()

Functional

functional

relu

relu(x)

Rectified linear unit, max(0, x).

Source code in minitorch/functional.py

def relu(x):
    """Rectified linear unit, `max(0, x)`."""
    data = np.maximum(0, x.data)
    out = Tensor(data, requires_grad=x.requires_grad)

    def _backward():
        if x.requires_grad:
            _accum_grad(x, (x.data > 0) * out.grad)

    out._backward = _backward
    out._prev = {x}
    out._op = 'relu'
    return out

gelu

gelu(x)

Gaussian error linear unit (tanh approximation, as in GPT-2).

Source code in minitorch/functional.py

def gelu(x):
    """Gaussian error linear unit (tanh approximation, as in GPT-2)."""
    # built from tensor ops so autograd differentiates it for free
    c = math.sqrt(2.0 / math.pi)
    inner = (x + x ** 3 * 0.044715) * c
    return x * 0.5 * (tanh(inner) + 1.0)

sigmoid

sigmoid(x)

Logistic sigmoid, 1 / (1 + e**-x).

Source code in minitorch/functional.py

def sigmoid(x):
    """Logistic sigmoid, `1 / (1 + e**-x)`."""
    s = 1.0 / (1.0 + np.exp(-x.data))
    out = Tensor(s, requires_grad=x.requires_grad)

    def _backward():
        if x.requires_grad:
            _accum_grad(x, out.data * (1.0 - out.data) * out.grad)

    out._backward = _backward
    out._prev = {x}
    out._op = 'sigmoid'
    return out

tanh

tanh(x)

Hyperbolic tangent.

Source code in minitorch/functional.py

def tanh(x):
    """Hyperbolic tangent."""
    t = np.tanh(x.data)
    out = Tensor(t, requires_grad=x.requires_grad)

    def _backward():
        if x.requires_grad:
            _accum_grad(x, (1.0 - out.data ** 2) * out.grad)

    out._backward = _backward
    out._prev = {x}
    out._op = 'tanh'
    return out

softmax

softmax(x, axis=-1)

Softmax over axis. Numerically stable via max-subtraction.

Source code in minitorch/functional.py

def softmax(x, axis=-1):
    """Softmax over `axis`. Numerically stable via max-subtraction."""
    shifted = x.data - x.data.max(axis=axis, keepdims=True)
    exps = np.exp(shifted)
    s = exps / exps.sum(axis=axis, keepdims=True)
    out = Tensor(s, requires_grad=x.requires_grad)

    def _backward():
        if x.requires_grad:
            g = out.grad * out.data
            g = g - out.data * g.sum(axis=axis, keepdims=True)
            _accum_grad(x, g)

    out._backward = _backward
    out._prev = {x}
    out._op = 'softmax'
    return out

log_softmax

log_softmax(x, axis=-1)

Log of softmax over axis. More stable than softmax(x).log().

Source code in minitorch/functional.py

def log_softmax(x, axis=-1):
    """Log of softmax over `axis`. More stable than `softmax(x).log()`."""
    shifted = x.data - x.data.max(axis=axis, keepdims=True)
    log_sum_exp = np.log(np.exp(shifted).sum(axis=axis, keepdims=True))
    data = shifted - log_sum_exp
    out = Tensor(data, requires_grad=x.requires_grad)

    def _backward():
        if x.requires_grad:
            s = np.exp(out.data)
            grad = out.grad - s * out.grad.sum(axis=axis, keepdims=True)
            _accum_grad(x, grad)

    out._backward = _backward
    out._prev = {x}
    out._op = 'log_softmax'
    return out