Un eje de lote cuesta una letra

gray       = np.einsum('hwc,c->hw',   photo, w)   # (512, 512)
gray_batch = np.einsum('nhwc,c->nhw', batch, w)   # (2, 512, 512)

Añadir un eje de lote cuesta exactamente una letra. La regla no cambia, la contracción no cambia, y la misma expresión funciona para una imagen o para un millón.

Por eso vale la pena aprender einsum: se lee como las matemáticas del capítulo 2 y escala sin reescribirse.