Skip to content

English · Español

Lab 01 — Ops elementwise con backward correcto para broadcasting

Objetivo: añadir la familia de ops elementwise a Tensor: add sub mul div neg exp log relu gelu tanh. Cada forward respeta el broadcasting de NumPy; cada backward suma correctamente a lo largo de los ejes broadcasteados. Cross-check de cada op contra PyTorch FP64 (tolerancia 1e-7) y gradcheck (tolerancia 1e-4 en FP64).

Tiempo estimado: 3–4 horas (son muchas ops, con una idea peliaguda — el broadcasting reverso).

Prerrequisitos: Lab 00. Teoría 02 (02-tensor-op-derivatives.md) — releer la derivación del broadcasting reverso.

Lo que produces

  • Las 9 ops elementwise en src/minitorch/tensor.py, expuestas vía dunders y métodos.
  • tests/test_elementwise.py con dos clases de tests por op: cross-check PyTorch + gradcheck.
  • tests/test_broadcast_reverse.py con los tests de pares broadcasting (un (3,) op un (4,3), un (N,1) op un (1,M), etc.).

El helper de broadcasting reverso

Según theory/02, el helper:

def _unbroadcast(grad: FloatArray, target_shape: tuple[int, ...]) -> FloatArray:
    # If grad.shape is "wider" than target_shape on the left, sum out the extra leading dims.
    while grad.ndim > len(target_shape):
        grad = grad.sum(axis=0)
    # For each remaining axis, if target dim is 1 but grad dim > 1, sum that axis with keepdims=True.
    for i, t_dim in enumerate(target_shape):
        if t_dim == 1 and grad.shape[i] > 1:
            grad = grad.sum(axis=i, keepdims=True)
    return grad

Implementa esto una vez. El _backward de cada op elementwise lo llama antes de asignar al .grad de un padre.

La trampa más común del lab: olvidar _unbroadcast y dejar un gradiente con forma (shape) distinta a la del parent. Al sumar parent.grad += new_grad NumPy difunde la suma silenciosamente. El resultado: una matriz donde esperabas un vector. Asegúrate de que tras cada _backward, parent.grad.shape == parent.data.shape.

TODOs (por op)

Para cada op, haz las mismas cinco cosas:

  1. Forward: calcula out_data con NumPy. Construye out = Tensor(out_data, _prev=(self, other), _op='<name>') y propaga requires_grad.
  2. Backward: define un closure capturando los padres y cualquier valor cacheado; al ser llamado, contribuye _unbroadcast(local_grad, parent.shape) al .grad de cada padre.
  3. Test PyTorch: construye la misma expresión en PyTorch FP64; compara gradientes con np.testing.assert_allclose(..., rtol=1e-7).
  4. Test gradcheck: FD numérico en FP64; tolerancia 1e-4.
  5. Test broadcasting: al menos un par de formas (shapes) donde el forward deba broadcastear.

Bloque A — familia aditiva

  • __add__(self, other): forward self.data + other.data. Backward: self.grad += _unbroadcast(out.grad, self.shape); lo mismo para other (trata escalares/ints de Python envolviéndolos en Tensor).
  • __radd__: para 2 + t.
  • __neg__(self): forward -self.data. Backward: self.grad += _unbroadcast(-out.grad, self.shape).
  • __sub__(self, other): reutiliza add + neg si prefieres; o implementa directamente. Documenta la elección.
  • __rsub__: para 2 - t.

Bloque B — familia multiplicativa

  • __mul__(self, other): forward self.data * other.data. Backward: self.grad += _unbroadcast(other.data * out.grad, self.shape); other.grad += _unbroadcast(self.data * out.grad, other.shape). Nota el swap: ∂(a*b)/∂a = b.
  • __rmul__.
  • __truediv__(self, other): forward self.data / other.data. Backward: ∂(a/b)/∂a = 1/b, ∂(a/b)/∂b = -a/b². Cachea other.data (no cachees 1/other.data — riesgo de división por cero en la caché).
  • __rtruediv__: para 2.0 / t.

Bloque C — no-linealidades unarias (métodos de Tensor)

  • exp(self): forward np.exp(self.data). Backward: self.grad += _unbroadcast(out.data * out.grad, self.shape) — nota que usamos out.data (que es igual a exp(self.data)).
  • log(self): forward np.log(self.data). Backward: self.grad += _unbroadcast(out.grad / self.data, self.shape). Asegura self.data > 0 en el forward (lanza ruidosamente si el input es inválido).
  • relu(self): forward np.maximum(0, self.data). Backward: self.grad += _unbroadcast(out.grad * (self.data > 0).astype(self.data.dtype), self.shape). Sub-gradiente en 0 = 0.
  • tanh(self): forward np.tanh(self.data). Backward: self.grad += _unbroadcast(out.grad * (1 - out.data**2), self.shape).
  • gelu(self): forward 0.5 * self.data * (1 + np.tanh(...)) (usa la aproximación tanh de la teoría). Backward: la derivada no es trivial — vuelve a derivarla en papel antes de programar; cachea intermedios.

Restricciones

  • Solo funcional. Sin ops in-place.
  • Sin numpy.errstate silenciando warnings. Deja que log(<=0) lance excepción.
  • Cada op recibe una línea de resumen en español en el eventual solutions/01-elementwise-ops-ref.md — mantén abierto tu theory/02-tensor-op-derivatives.md para las derivaciones.
  • Usa las utilidades de seeding de la fase 6 en los tests: seed_everything(42).

Patrones de test

Cross-check PyTorch (un ejemplo)

import torch, numpy as np
from minitorch import Tensor

def test_mul_pytorch_cross():
    rng = np.random.default_rng(42)
    a_data = rng.standard_normal((3, 4)).astype(np.float64)
    b_data = rng.standard_normal((3, 4)).astype(np.float64)
    # Our autograd.
    A = Tensor(a_data, requires_grad=True)
    B = Tensor(b_data, requires_grad=True)
    L = (A * B).sum()
    L.backward()
    # PyTorch reference.
    At = torch.tensor(a_data, dtype=torch.float64, requires_grad=True)
    Bt = torch.tensor(b_data, dtype=torch.float64, requires_grad=True)
    Lt = (At * Bt).sum()
    Lt.backward()
    np.testing.assert_allclose(A.grad, At.grad.numpy(), rtol=1e-7)
    np.testing.assert_allclose(B.grad, Bt.grad.numpy(), rtol=1e-7)

Escribe uno de estos por op. ~10 líneas cada uno.

Gradcheck (un ejemplo)

def test_mul_gradcheck():
    rng = np.random.default_rng(0)
    a = Tensor(rng.standard_normal((2, 3)), requires_grad=True)
    b = Tensor(rng.standard_normal((2, 3)), requires_grad=True)
    f = lambda: (a * b).sum()
    assert gradcheck_pair(f, a, b, eps=1e-6, atol=1e-4)

Test de broadcasting (un ejemplo, con sabor gramatical)

def test_mul_broadcasts_person_logits():
    # (3,) person one-hot times (3, 5) (person, tense) logits.
    person = Tensor(np.array([0.0, 1.0, 0.0]), requires_grad=True)
    logits = Tensor(np.random.default_rng(7).standard_normal((3, 5)), requires_grad=True)
    out = (person[:, None] * logits).sum()
    out.backward()
    assert person.grad.shape == (3,)        # _unbroadcast restored
    assert logits.grad.shape == (3, 5)

Condiciones de parada

Hecho cuando:

  1. Las 9 ops elementwise implementadas.
  2. Cada una tiene al menos 2 tests: cross PyTorch + gradcheck. Todos verdes.
  3. Cada una tiene al menos 1 test de broadcasting. Todos verdes.
  4. mypy --strict limpio.
  5. Puedes responder: "si olvido _unbroadcast en el backward de __add__, ¿cuál es el test más pequeño que lo detecta?"

Escollos

  • out.grad es None en la hoja. La primera llamada a _backward es root._backward() tras root.grad = np.array(1.0). Si escribiste _backward para leer out.grad, bien. Si por accidente capturaste out y luego lees out.grad antes de que fuera fijado por el paso reverse anterior, obtienes None. Traza el orden con cuidado.
  • Signo erróneo en __rsub__ / __rtruediv__. 2 - t es (-t) + 2, no t - 2. Fácil de invertir.
  • Cachear 1/other.data explota. Mantén other.data y divide en el closure; o protégete contra el cero.
  • Gradiente de relu exactamente en 0. Convención: 0. No uses accidentalmente (self.data >= 0) — esa es una convención distinta.
  • El forward y el backward de gelu no coinciden. Dos formulaciones de GELU (exacta vía erf y aproximación tanh). Escoge una y úsala consistentemente.

Cuándo consultar solutions/

Después de que las 9 ops pasen las tres familias de tests. solutions/01-elementwise-ops-ref.md (al abrir la fase) compara tus closures contra las implementaciones canónicas.

Siguiente lab: lab/02-reduction-and-shape-ops.md.