Skip to content

English · Español

Lab 00 — Predecir shapes a partir de strings einsum

Objetivo: hacer mecánica la aritmética de shapes con einsum. Sin código en la Parte A; solo papel. Los ejemplos están anclados en las codificaciones de formas verbales §A13.

Tiempo estimado: 60–90 minutos.

Prerrequisito: teoría 01-tensors-and-shapes.md y 02-matmul-and-shapes.md leídas.

Lo que produces

Un directorio experiments/03-shapes-by-hand/ que contiene:

  • predictions.md — tus predicciones de shape escritas a mano para cada problema de abajo.
  • verify.py — script corto que construye los operandos y ejecuta los einsums, comparando con tus predicciones.
  • results.json — pass/fail por predicción.
  • manifest.json.

Las constantes de dimensión §A13

Para cada problema, asume estos tamaños estándar (según theory/01-tensors-and-shapes.md):

B   = 32      # batch
T   = 16      # longitud de secuencia
V   = 600     # vocabulario (§A13: 20 verbos × 5 tiempos × 3 personas + pares en español)
D   = 64      # dim de embedding
H   = 4       # número de cabezas de attention
D_k = 16      # dim por cabeza = D / H
D_ff= 256     # intermedia del FFN
K_classes = 5 # número de clases de tiempo verbal

TODOs

Parte A — predecir en papel

Para cada expresión einsum de abajo, escribe en predictions.md:

  • El shape de cada operando.
  • El shape de la salida.
  • El número total de FLOPs multiply-add (2 × producto_de_todos_los_índices).
  • Una descripción de una frase en inglés de lo que la operación hace en el contexto §A13.

Sin código todavía. Resuelve leyendo el string einsum y aplicando las dos reglas (repetido = suma, libre = salida).

1. Búsqueda de embedding (un solo token).

einsum('v,vd->d', one_hot, E)

one_hot.shape = (V,), E.shape = (V, D).

2. Búsqueda de embedding por lotes.

einsum('btv,vd->btd', tokens_one_hot, E)

tokens_one_hot.shape = (B, T, V), E.shape = (V, D).

3. Clasificación de tiempo verbal.

einsum('bd,kd->bk', hidden, W_tense)

hidden.shape = (B, D), W_tense.shape = (K_classes, D).

4. Clasificación de tiempo verbal por token (con batch + secuencial).

einsum('btd,kd->btk', x, W_tense)

x.shape = (B, T, D), W_tense.shape = (K_classes, D).

5. Proyección lineal (Q en attention).

einsum('btd,de->bte', x, W_Q)

x.shape = (B, T, D), W_Q.shape = (D, D).

6. Reshape para multi-head — split de Q, K, V.

Tras calcular Q de shape (B, T, D), haces reshape a (B, T, H, D_k) y luego transpose a (B, H, T, D_k). Escribe el einsum que va de (B, T, D) a (B, H, T, D_k) directamente. (Pista: piénsalo como una contracción no-op con un reshape apropiado.) En realidad einsum no puede hacer reshapes por sí solo; en su lugar, predice los shapes tras Q.reshape(B, T, H, D_k).transpose(0, 2, 1, 3).

7. Scores de attention.

einsum('bhqd,bhkd->bhqk', Q, K)

Q.shape = (B, H, T, D_k), K.shape = (B, H, T, D_k).

8. Salida de attention.

einsum('bhqk,bhkd->bhqd', attn_probs, V)

attn_probs.shape = (B, H, T, T), V.shape = (B, H, T, D_k).

9. Proyección de salida.

einsum('bhtd,hde->bte', attn_out, W_O)

attn_out.shape = (B, H, T, D_k), W_O.shape = (H, D_k, D).

10. Expansión FFN.

einsum('btd,df->btf', x, W_1)

x.shape = (B, T, D), W_1.shape = (D, D_ff).

11. Contracción FFN.

einsum('btf,fd->btd', h, W_2)

h.shape = (B, T, D_ff), W_2.shape = (D_ff, D).

12. Proyección al vocabulario (capa final).

einsum('btd,vd->btv', x, E)

x.shape = (B, T, D), E.shape = (V, D). (Nota que E aquí está atada al embedding de entrada — la misma matriz.)

13. Reducción del log-likelihood de entropía cruzada.

einsum('btv,btv->', log_probs, labels_one_hot)

log_probs.shape = (B, T, V), labels_one_hot.shape = (B, T, V).

14. Log-likelihood medio por secuencia.

einsum('btv,btv->b', log_probs, labels_one_hot)

(Luego divide por T.) ¿Cuál es el shape resultante?

15. Diagonal de una matriz cuadrada.

einsum('ii->i', M)

M.shape = (5, 5).

16. Traza.

einsum('ii->', M)

M.shape = (5, 5).

17. Producto interno de Frobenius.

einsum('mn,mn->', A, B)

A.shape = (20, 15), B.shape = (20, 15).

18. Producto externo de dos vectores §A13 de formas verbales.

einsum('v,w->vw', a, b)

a.shape = (V,), b.shape = (V,). ¿Cuál es el tamaño en MB en fp32?

19. Contracción mixta con batch.

einsum('btd,dvf->btvf', x, T)

x.shape = (B, T, D), T.shape = (D, V, D_ff). (Inusual; solo para practicar shapes.)

20. El dot de la matriz de conteos de conjugaciones §A13.

einsum('vp,wp->vw', C, C)

C.shape = (20, 15) (20 verbos × 15 índices de conjugación). ¿Qué está calculando C @ C^T en términos §A13?

Parte B — verificar con código

verify.py: para cada una de las 20 expresiones de arriba, construye operandos aleatorios con los shapes especificados (usa np.random.default_rng(42).standard_normal(shape).astype(np.float32)), ejecuta el einsum, imprime el shape real y compáralo con tu predicción.

predictions = {
    1: (D,),
    2: (B, T, D),
    3: (B, K_classes),
    # ...
}

for expr_id, expected_shape in predictions.items():
    # construir operandos
    # ejecutar einsum
    # comprobar shape
    pass_fail = (actual_shape == expected_shape)
    print(f"{expr_id}: predicted {expected_shape}, got {actual_shape}, {'PASS' if pass_fail else 'FAIL'}")

Guarda los resultados en results.json. Debes conseguir 20/20. Un fallo significa volver a derivar en papel antes de volver a ejecutar.

Parte C — verificación de FLOPs

Para tres de las expresiones (tu elección — prueba 2, 7, 12), calcula los FLOPs teóricos (escribe la fórmula en predictions.md). Compara con el tiempo medido × los GFLOPS pico de tu máquina (del roofline de la Fase 1). El tiempo medido puede ser mucho mayor que lo que predice la teoría por la sobrecarga de Python.

Parte D — la pregunta asesina

La expresión 12 ('btd,vd->btv') y la expresión 2 ('btv,vd->btd') parecen casi inversas la una de la otra. ¿Lo son? En términos §A13, la expresión 2 es "busca el embedding para cada token"; la expresión 12 es "proyecta el estado oculto de vuelta a logits del vocabulario". ¿Qué propiedad de la matriz de embedding E haría que fueran inversas genuinas? (Pista: tiene relación con la ortogonalidad.) Discútelo en predictions.md.

Restricciones

  • Predice antes de ejecutar. Parte A → Parte B, en orden.
  • Descripción de una frase en inglés para cada una. Fuerza a pensar en términos §A13, no solo en aritmética de shapes.
  • np.einsum es la única implementación permitida — incluso para operaciones que tienen funciones NumPy especializadas (np.dot, np.matmul). El objetivo es practicar einsum.

Condiciones de parada

Hecho cuando:

  1. predictions.md tiene las 20 predicciones con shape + FLOPs + descripción en inglés.
  2. verify.py imprime 20/20 PASS.
  3. La pregunta asesina de la Parte D tiene una respuesta escrita.
  4. Puedes leer cualquier nuevo string einsum y predecir su shape sin consultar notas.

Escollos

  • Índice faltante en la salida. Si escribes 'btv,vd->btv', la d queda sin emparejar en la salida — einsum inválido. NumPy lanza un error.
  • Dimensión inconsistente en distintos operandos. Si tokens_one_hot.shape = (B, T, 599) en lugar de (B, T, V=600), la v en E (tamaño 600) no coincidirá. NumPy lanza un error.
  • El -> final. Si omites el lado derecho, numpy usa una convención implícita (suma cualquier eje que no aparezca en el lado derecho, manteniendo orden alfabético para no repetidos). Sé explícito siempre.

Cuándo consultar solutions/

Tras commitear los cuatro archivos. Solución en solutions/00-shapes-by-hand-ref.md (escrita en la apertura de fase).

Siguiente lab: lab/01-matmul-perf.md.