English · Español

Lab 02 — Ops de reducción y de forma (shape)¶

Objetivo: añadir la familia de reducción + forma (shape) a Tensor: sum mean reshape transpose broadcast_to getitem cat stack. Las ops de "forma" son no-ops en sentido matemático (Jacobiano local es identidad) pero requieren contabilidad cuidadosa de índices en el backward. Las ops de "reducción" son lo contrario: la matemática es trivial, pero el backward debe broadcast_to el gradiente upstream de vuelta a la forma (shape) de entrada.

Tiempo estimado: 3–4 horas.

Prerrequisitos: Labs 00 + 01. Teoría 02-tensor-op-derivatives.md releída (especialmente la derivación del backward de reducción).

Lo que produces¶

Las 7 ops añadidas a src/minitorch/tensor.py.
tests/test_reductions.py con cross-check + gradcheck por op de reducción.
tests/test_shape_ops.py con tests de roundtrip de forma (shape) + gradcheck por op de forma.

Recordatorio del contrato de forma (shape)¶

Después de cada _backward, parent.grad.shape == parent.data.shape. Para ops de reducción, esto significa broadcastear el gradiente upstream de vuelta a la forma (shape) de entrada. Para ops de forma, esto significa invertir la transformación de forma sobre el gradiente.

La regla mental: una reduction colapsa ejes hacia adelante, y los expande hacia atrás; una reshape reorganiza el layout, y el backward solo desreorganiza. No hay álgebra nueva, solo cuidar el shape.

TODOs (por op)¶

Para cada op de abajo, haz las cinco cosas del Lab 01 (forward, backward, test PyTorch, test gradcheck, test de borde de forma (shape)). Las ops de forma no necesitan tests de broadcasting.

Bloque A — reducciones¶

sum(self, axis=None, keepdims=False): forward self.data.sum(axis=..., keepdims=...). Backward: el gradiente upstream tiene forma (shape) igual a la salida del forward; debes hacer broadcast_to(self.shape) después de potencialmente reinsertar los ejes reducidos.
Pista de implementación: si keepdims=False y axis no es None, el out.grad upstream tiene rango self.data.ndim - len(axes). Usa np.expand_dims en cada eje reducido para restaurar las dims reducidas como tamaño-1, luego np.broadcast_to(self.shape).
Cuando axis=None, la salida es un escalar; el backward es np.broadcast_to(out.grad, self.shape).
mean(self, axis=None, keepdims=False): forward self.data.mean(axis=..., keepdims=...). Backward idéntico a sum pero divide por N, donde N es el número de elementos reducidos. Cachea N en el closure.

Bloque B — ops de forma (shape) puras¶

reshape(self, new_shape: tuple[int, ...]): forward self.data.reshape(new_shape). Backward self.grad += out.grad.reshape(self.shape). Sin cambio de tamaño, sin matemáticas.
transpose(self, axes: tuple[int, ...] | None = None): forward self.data.transpose(axes). Backward: transpón el upstream por la permutación inversa. Pista: si axes = (1, 0, 2), la inversa es argsort(axes) = (1, 0, 2) (involución aquí); en general usa np.argsort(axes).
broadcast_to(self, shape: tuple[int, ...]): forward np.broadcast_to(self.data, shape) (devuelve una vista (view); si tus ops aguas abajo mutan, haz una copia — pero no mutan, así que la vista (view) sirve). Backward: esta es la única op de forma (shape) que necesita _unbroadcast — el gradiente vuelve con shape y debe reducirse a self.shape.

Bloque C — indexación/concat¶

__getitem__(self, key): forward self.data[key]. Backward: reserva np.zeros_like(self.data), dispersa el gradiente upstream en key. Pista: grad[key] += out.grad funciona para slices básicos e indexación fancy sin índices duplicados. Si key pudiera contener duplicados (p. ej., self[[0, 0, 1]]), usa np.add.at(grad, key, out.grad) para asegurar la acumulación.
cat(tensors: list[Tensor], axis: int) (función a nivel de módulo, no un método): forward np.concatenate([t.data for t in tensors], axis=axis). Backward: troceá el gradiente upstream a lo largo de axis según el tamaño de cada tensor en ese eje; contribuye cada trozo al padre correspondiente.
stack(tensors: list[Tensor], axis: int): forward np.stack([t.data for t in tensors], axis=axis). Backward: divide a lo largo del nuevo eje y haz squeeze.

Casos peliagudos a cubrir¶

Caso 1 — `sum` con `keepdims=False` sobre un único eje¶

x = Tensor(np.random.randn(3, 4), requires_grad=True)
y = x.sum(axis=0)      # shape (4,)
y.sum().backward()
assert x.grad.shape == (3, 4)
assert np.allclose(x.grad, 1.0)  # uniform gradient

Si olvidas reinsertar axis=0 antes del broadcast, el out.grad de forma (shape) (4,) broadcastea a (3, 4) solo porque NumPy lo alinea a la derecha por defecto — pero para axis=1 (salida de forma (shape) (3,)), la alineación sería incorrecta y obtendrías el gradiente equivocado.

Caso 2 — `mean` sobre `axis=None`¶

x = Tensor(np.random.randn(2, 3), requires_grad=True)
y = x.mean()
y.backward()
assert np.allclose(x.grad, 1.0 / 6)  # uniform 1/N gradient

Caso 3 — `getitem` con índices duplicados¶

x = Tensor(np.array([1.0, 2.0, 3.0]), requires_grad=True)
y = x[[0, 0, 1]]   # shape (3,) selecting x[0] twice and x[1] once
y.sum().backward()
assert np.allclose(x.grad, [2.0, 1.0, 0.0])  # x[0] contributes twice

Si escribiste x.grad[key] += out.grad, obtendrías [1.0, 1.0, 0.0] (el duplicado se sobreescribe, no se acumula). Usa np.add.at.

Caso 4 — `transpose` de un tensor rank-3 con una permutación no trivial¶

x = Tensor(np.random.randn(2, 3, 4), requires_grad=True)
y = x.transpose((2, 0, 1))   # shape (4, 2, 3)
y.sum().backward()
assert x.grad.shape == (2, 3, 4)

La permutación inversa: np.argsort((2, 0, 1)) = (1, 2, 0). El backward aplica out.grad.transpose((1, 2, 0)) para recuperar la forma (shape) (2, 3, 4).

Caso 5 — `cat` de dos tensores con tamaños diferentes en el eje cat¶

a = Tensor(np.random.randn(2, 3), requires_grad=True)
b = Tensor(np.random.randn(2, 5), requires_grad=True)
y = cat([a, b], axis=1)      # shape (2, 8)
y.sum().backward()
assert a.grad.shape == (2, 3)
assert b.grad.shape == (2, 5)

El backward debe recordar el tamaño por tensor a lo largo del eje cat para trocear correctamente. Cachea [t.shape[axis] for t in tensors] en el closure.

Restricciones¶

Sin nuevas deps externas. NumPy + librería estándar únicamente.
Sin ops in-place. reshape devuelve siempre un nuevo Tensor (su data puede ser una vista (view), pero la identidad del Tensor es nueva — importante para el DAG).
broadcast_to devuelve una vista (view) de self.data. Documéntalo. El forward comparte memoria con el padre; el backward re-reserva vía _unbroadcast. Este es el patrón que PyTorch usa con expand (frente a repeat, que copia).

Patrones de test¶

Cross-check de reducción (un ejemplo)¶

def test_sum_axis_pytorch_cross():
    rng = np.random.default_rng(42)
    a = rng.standard_normal((4, 5)).astype(np.float64)
    A = Tensor(a, requires_grad=True)
    L = A.sum(axis=0).sum()  # scalar
    L.backward()
    At = torch.tensor(a, dtype=torch.float64, requires_grad=True)
    Lt = At.sum(dim=0).sum()
    Lt.backward()
    np.testing.assert_allclose(A.grad, At.grad.numpy(), rtol=1e-7)

Gradcheck de forma (shape) (un ejemplo, con sabor gramatical)¶

def test_reshape_grammar_grid_gradcheck():
    # (3, 5) person-tense logits flattened to (15,) and back.
    rng = np.random.default_rng(0)
    grid = Tensor(rng.standard_normal((3, 5)), requires_grad=True)
    f = lambda: grid.reshape((15,)).sum()
    assert gradcheck(f, grid, eps=1e-6, atol=1e-4)

Test de indexación (con sabor gramatical)¶

def test_getitem_select_third_person():
    # Pick out third-person row from (3, 5) logits.
    logits = Tensor(np.random.default_rng(0).standard_normal((3, 5)), requires_grad=True)
    out = logits[2].sum()         # third person, all tenses
    out.backward()
    expected = np.zeros((3, 5))
    expected[2] = 1.0
    np.testing.assert_allclose(logits.grad, expected)

Condiciones de parada¶

Hecho cuando:

Las 7 ops implementadas.
Cada una tiene ≥ 2 tests (cross PyTorch + gradcheck). Todos verdes.
Los casos peliagudos 1–5 todos verdes.
mypy --strict limpio. cat y stack están tipados como Callable[[list[Tensor], int], Tensor] a nivel de módulo.
Puedes responder: "¿por qué el backward de broadcast_to necesita _unbroadcast pero el de reshape no?"

Escollos¶

Backward de reducción con keepdims=False. El gradiente upstream está reducido en rango; debes reinsertar los ejes reducidos como tamaño-1 antes de broadcastear. Olvidar esto da un gradiente silenciosamente incorrecto (dirección de suma errónea).
Permutación inversa en transpose. np.argsort(axes) es la inversa correcta para cualquier permutación. Derivar la inversa a mano es propenso a errores — usa argsort.
getitem con máscaras booleanas. Usa np.add.at para el backward; asignación-con-máscara funciona en forward pero es frágil en backward.
Identidad de padre en cat/stack. Si el mismo tensor aparece dos veces en la lista de entrada (cat([a, a], axis=0)), a.grad debe acumularse desde ambos trozos. Testéalo — el bug es silencioso.
broadcast_to es una vista (view); repeat es una copia. No los confundas. Implementamos broadcast_to porque toda op elementwise ya broadcastea implícitamente; repeat es innecesario en la fase 8.

Cuándo consultar `solutions/`¶

Después de que las 7 ops pasen todos los tests. solutions/02-reduction-and-shape-ops-ref.md comparará tus closures de _unbroadcast/reshape contra las implementaciones canónicas.

Siguiente lab: lab/03-matmul-softmax-ce.md.