Broadcasting

Spec-level broadcasting and broadcasted binary maps.

Broadcasting is defined in terms of the shape proposition Shape.CanBroadcastTo: pick an explicit target shape and provide evidence that each operand broadcasts to it. The proposition depends only on the two shapes, so the result never depends on how the evidence was obtained.

Execution goes through the internal Rep.broadcast: the source is padded with leading singleton axes by a zero-copy reshape and then pulled back along the coordinate map that fixes expanded axes at index zero. This is one buffer fill, independent of the rank.

TorchLean standardizes on this explicit target style throughout core. There is intentionally no second "implicit" API that infers a common output shape from two operands, because that would split the codebase into two parallel styles. The backward pass of a broadcast is the sum-reduction over the broadcast axes (reduceFromBroadcastTo in the reductions module).

theorem Spec.Shape.CanBroadcastTo.forall₂_toList {s₁ s₂ : Shape} (h : s₁.CanBroadcastTo s₂) : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (padLeft (s₂.rank - s₁.rank) s₁).toList s₂.toList

Shape.CanBroadcastTo in the list form used by the internal representation: after padding the source with leading ones to the target rank, extents agree pointwise or the source extent is one.

def TorchLean.Tensor.Broadcasting.Internal.broadcastPadded {α : Type} [Storage α] {s₁ s₂ : Spec.Shape} (k : ℕ) (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (Spec.Shape.padLeft k s₁).toList s₂.toList) (tensor : Tensor α s₁) : Tensor α s₂

Broadcast a source padded with k leading singleton axes.

Separating the padding count from the shapes lets the structural lemmas below be stated for literal counts 0 and k + 1, independently of the rank arithmetic in broadcastTo.

theorem TorchLean.Tensor.Broadcasting.Internal.broadcastPadded_congr {α : Type} [Storage α] {s₁ s₂ : Spec.Shape} {k k' : ℕ} (hk : k = k') (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (Spec.Shape.padLeft k s₁).toList s₂.toList) (tensor : Tensor α s₁) : broadcastPadded k h tensor = broadcastPadded k' ⋯ tensor

Transport the padding count along an equality.

theorem TorchLean.Tensor.Broadcasting.Internal.broadcastPadded_zero_scalar {α : Type} [Storage α] (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (Spec.Shape.padLeft 0 Spec.Shape.scalar).toList Spec.Shape.scalar.toList) (tensor : Tensor α Spec.Shape.scalar) : broadcastPadded 0 h tensor = tensor

Broadcasting a scalar to the scalar shape is the identity.

theorem TorchLean.Tensor.Broadcasting.Internal.broadcastPadded_succ {α : Type} [Storage α] {k n : ℕ} {s t : Spec.Shape} (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (Spec.Shape.padLeft (k + 1) s).toList (Spec.Shape.dim n t).toList) (tensor : Tensor α s) : broadcastPadded (k + 1) h tensor = dim fun (x : Fin n) => broadcastPadded k ⋯ tensor

One more padding axis stacks the broadcast of the remaining padding along the leading target axis.

theorem TorchLean.Tensor.Broadcasting.Internal.broadcastPadded_zero_dim_eq {α : Type} [Storage α] {n : ℕ} {s t : Spec.Shape} (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (Spec.Shape.padLeft 0 (Spec.Shape.dim n s)).toList (Spec.Shape.dim n t).toList) (tensor : Tensor α (Spec.Shape.dim n s)) : broadcastPadded 0 h tensor = dim fun (i : Fin n) => broadcastPadded 0 ⋯ (tensor.unstack i)

With no padding and equal leading extents, broadcasting acts slice by slice.

theorem TorchLean.Tensor.Broadcasting.Internal.broadcastPadded_zero_dim_one {α : Type} [Storage α] {n : ℕ} {s t : Spec.Shape} (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (Spec.Shape.padLeft 0 (Spec.Shape.dim 1 s)).toList (Spec.Shape.dim n t).toList) (tensor : Tensor α (Spec.Shape.dim 1 s)) : broadcastPadded 0 h tensor = dim fun (x : Fin n) => broadcastPadded 0 ⋯ (tensor.unstack 0)

With no padding and a leading source extent of one, every target slice reads slice zero.

def TorchLean.Tensor.broadcastTo {α : Type} [Storage α] {s₁ s₂ : Spec.Shape} (h : s₁.CanBroadcastTo s₂) (tensor : Tensor α s₁) : Tensor α s₂

Broadcast a tensor to a compatible target shape (spec-level analogue of torch.broadcast_to).

The source is padded with leading singleton axes to the target rank and every singleton axis is replicated. The result depends only on the shapes, not on how h was proved.

@[simp]

theorem TorchLean.Tensor.broadcastTo_scalar {α : Type} [Storage α] (h : Spec.Shape.scalar.CanBroadcastTo Spec.Shape.scalar) (tensor : Tensor α Spec.Shape.scalar) : broadcastTo h tensor = tensor

Broadcasting a scalar tensor to the scalar shape is the identity.

theorem TorchLean.Tensor.broadcastTo_dim_eq {α : Type} [Storage α] {n : ℕ} {s t : Spec.Shape} (hRank : s.rank = t.rank) (h : (Spec.Shape.dim n s).CanBroadcastTo (Spec.Shape.dim n t)) (tensor : Tensor α (Spec.Shape.dim n s)) : broadcastTo h tensor = dim fun (i : Fin n) => broadcastTo ⋯ (tensor.unstack i)

Equal leading extents broadcast slice by slice.

theorem TorchLean.Tensor.broadcastTo_dim_one {α : Type} [Storage α] {n : ℕ} {s t : Spec.Shape} (hRank : s.rank = t.rank) (h : (Spec.Shape.dim 1 s).CanBroadcastTo (Spec.Shape.dim n t)) (tensor : Tensor α (Spec.Shape.dim 1 s)) : broadcastTo h tensor = dim fun (x : Fin n) => broadcastTo ⋯ (tensor.unstack 0)

A leading source extent of one is replicated along the leading target axis.

theorem TorchLean.Tensor.broadcastTo_expand {α : Type} [Storage α] {n : ℕ} {s t : Spec.Shape} (hRank : s.rank ≤ t.rank) (h : s.CanBroadcastTo (Spec.Shape.dim n t)) (tensor : Tensor α s) : broadcastTo h tensor = dim fun (x : Fin n) => broadcastTo ⋯ tensor

A target axis that only raises the rank replicates the whole source.

@[simp]

theorem TorchLean.Tensor.broadcastTo_self {α : Type} [Storage α] {shape : Spec.Shape} (h : shape.CanBroadcastTo shape) (tensor : Tensor α shape) : broadcastTo h tensor = tensor

Broadcasting a tensor to its own shape changes nothing.

@[simp]

theorem TorchLean.Tensor.broadcastTo_dim_self {α : Type} [Storage α] {n : ℕ} {s : Spec.Shape} (h : s.CanBroadcastTo (Spec.Shape.dim n s)) (tensor : Tensor α s) : broadcastTo h tensor = dim fun (x : Fin n) => tensor

Broadcasting across one new leading target axis replicates the source.

Broadcasted maps

def TorchLean.Tensor.mapScalarLeft {α : Type} [Storage α] (f : α → α → α) (x : α) {s : Spec.Shape} (tensor : Tensor α s) : Tensor α s

Helper: map a scalar on the left over any tensor shape.

def TorchLean.Tensor.mapScalarRight {α : Type} [Storage α] (f : α → α → α) (y : α) {s : Spec.Shape} (tensor : Tensor α s) : Tensor α s

Helper: map a scalar on the right over any tensor shape.

def TorchLean.Tensor.broadcastMapTo {α : Type} [Storage α] (f : α → α → α) {s₁ s₂ t : Spec.Shape} (cbx : s₁.CanBroadcastTo t) (cby : s₂.CanBroadcastTo t) : Tensor α s₁ → Tensor α s₂ → Tensor α t

Binary element-wise operation with broadcasting to an explicit target shape.

This is the helper you typically want in spec code:

• pick the output shape t,
• broadcast each operand to t,
• then map2Spec the pointwise operation.

PyTorch analogy: f(x, y) where x and/or y are broadcastable to a common shape. The common shape is explicit rather than discovered at runtime, which makes the result type predictable and fixes the intended output shape at the call site.