BugZoo: shape and broadcasting boundaries

This file records the TorchLean response to a very common bug class: tensor shape mistakes that would normally appear only at runtime, or worse, silently broadcast to the wrong expression.

The empirical motivation is concrete. Wu, Shen, and Chen build SFData, a corpus of 146 crashing tensor-shape bugs, and use missing batch dimensions as a representative pattern:

• Wu, Shen, and Chen, “Detecting Tensor Shape Faults in Deep Learning Systems”, ISSTA 2022. https://doi.org/10.1145/3533767.3534383

Wang et al. also describe numerical bugs where a missing keep_dims=True changes a reduction shape; NumPy/PyTorch-style broadcasting then makes a later expression typecheck while computing the wrong loss:

• Wang et al., “An Empirical Study on Numerical Bugs in Deep Learning Programs”, ASE NIER 2022. https://doi.org/10.1145/3551349.3559561

TorchLean makes this case explicit: ordinary elementwise ops require the same shape, and broadcasting requires Shape.CanBroadcastTo evidence. The examples below show the intended workflow. A missing batch dimension is not silently accepted; we write the singleton batch insertion. A reduced vector is not silently expanded back into a matrix; we carry the broadcast proof.

Bug-shaped PyTorch sketch:

# Crashes or silently changes later code depending on where it appears:
image = torch.randn(100, 100, 3)
model(image)          # model expected [1, 100, 100, 3]

# Easy to miss: reduction drops a dimension, then a later op broadcasts it back.
row_sum = x.sum(dim=0)          # [3], not [1, 3]
loss = ((x - row_sum) ** 2).sum()

TorchLean equivalent:

def batched := addSingletonBatch image
def row := reduceRows x
def explicit := broadcastRowToMatrix row

The important part is not the syntax; it is that the shape change and broadcast are named terms with types and proof evidence.

@[reducible, inline]

abbrev NN.Examples.BugZoo.ShapeAndBroadcast.ImageShape : Spec.Shape

A single HWC image: 100 x 100 pixels, three channels.

@[reducible, inline]

abbrev NN.Examples.BugZoo.ShapeAndBroadcast.SingletonBatchImageShape : Spec.Shape

The same image with a batch axis of one in front.

def NN.Examples.BugZoo.ShapeAndBroadcast.addSingletonBatch {α : Type} [TorchLean.Storage α] (x : TorchLean.Tensor α ImageShape) : TorchLean.Tensor α SingletonBatchImageShape

Insert an explicit singleton batch dimension.

This is the TorchLean version of the fix for the classic “forgot the batch axis” bug: we do not let Tensor α [100,100,3] masquerade as Tensor α [1,100,100,3]; the user has to name the reshape.

@[simp]

theorem NN.Examples.BugZoo.ShapeAndBroadcast.addSingletonBatch_zero {α : Type} [TorchLean.Storage α] (x : TorchLean.Tensor α ImageShape) : (addSingletonBatch x)[0] = x

Reading the only batch entry after addSingletonBatch gives back the original image.

@[reducible, inline]

abbrev NN.Examples.BugZoo.ShapeAndBroadcast.MatrixShape : Spec.Shape

A small 2 x 3 matrix, used to show what a reduction does to the shape.

@[reducible, inline]

abbrev NN.Examples.BugZoo.ShapeAndBroadcast.RowShape : Spec.Shape

The row vector left after reducing MatrixShape over its outer axis.

def NN.Examples.BugZoo.ShapeAndBroadcast.reduceRows {α : Type} [TorchLean.Storage α] [Add α] [Zero α] (x : TorchLean.Tensor α MatrixShape) : TorchLean.Tensor α RowShape

Sum over the outer axis of a 2 × 3 tensor, dropping that axis and producing a row vector.

theorem NN.Examples.BugZoo.ShapeAndBroadcast.rowBroadcastToMatrix : RowShape.CanBroadcastTo MatrixShape

Evidence that a row vector can be broadcast back across the outer dimension of a 2 × 3 matrix.

This is exactly the piece TorchLean wants users and proof scripts to make visible: if a reduction dropped a dimension, any later expansion is an explicit broadcast, not an accidental side effect.

def NN.Examples.BugZoo.ShapeAndBroadcast.broadcastRowToMatrix {α : Type} [TorchLean.Storage α] [Inhabited α] (x : TorchLean.Tensor α RowShape) : TorchLean.Tensor α MatrixShape

Broadcast a row vector to every row of a 2 × 3 matrix, using the evidence above.

def NN.Examples.BugZoo.ShapeAndBroadcast.inferredRowBroadcastToMatrix {α : Type} [TorchLean.Storage α] [Inhabited α] (x : TorchLean.Tensor α RowShape) : TorchLean.Tensor α MatrixShape

The inferred broadcast follows NumPy/PyTorch's right-aligned convention.

@[simp]

theorem NN.Examples.BugZoo.ShapeAndBroadcast.broadcastRowToMatrix_firstRow {α : Type} [TorchLean.Storage α] [Inhabited α] (x : TorchLean.Tensor α RowShape) : (broadcastRowToMatrix x)[0] = x

The first row of an explicit broadcast is definitionally the original row.

theorem NN.Examples.BugZoo.ShapeAndBroadcast.inferredRowBroadcastToMatrix_eq {α : Type} [TorchLean.Storage α] [Inhabited α] (x : TorchLean.Tensor α RowShape) : inferredRowBroadcastToMatrix x = broadcastRowToMatrix x

Inference and the explicit right-aligned witness compute the same matrix.

The dimensions in the preceding example are different, so a reversed alignment convention would fail rather than compute a different tensor. The square example below is the sharper regression: both axes have length two, but NumPy/PyTorch broadcasting still requires the vector to align with the final axis.

@[reducible, inline]

abbrev NN.Examples.BugZoo.ShapeAndBroadcast.SquareRowShape : Spec.Shape

@[reducible, inline]

abbrev NN.Examples.BugZoo.ShapeAndBroadcast.SquareMatrixShape : Spec.Shape

A 2 x 2 matrix, the target of the broadcast below.

def NN.Examples.BugZoo.ShapeAndBroadcast.inferredSquareBroadcast {α : Type} [TorchLean.Storage α] [Inhabited α] (x : TorchLean.Tensor α SquareRowShape) : TorchLean.Tensor α SquareMatrixShape

Broadcast a length-two vector across the rows of a 2 x 2 matrix.

theorem NN.Examples.BugZoo.ShapeAndBroadcast.inferredSquareBroadcast_rows {α : Type} [TorchLean.Storage α] [Inhabited α] (x : TorchLean.Tensor α SquareRowShape) : (inferredSquareBroadcast x)[0] = x ∧ (inferredSquareBroadcast x)[1] = x

The inferred square broadcast is right-aligned: each matrix row is the source vector.