DAG Shape and Reduction Primitives

Reductions, transposes, gathers, broadcasting, concatenation, slicing, and shape-preserving views.

def NN.GraphSpec.DAG.PrimOp.reduceSum (s : Spec.Shape) (axis : ℕ) [Spec.Shape.HasNonemptyAxis axis s] [s.WellFormed] : PrimOp [s] (TorchLean.Tensor.shapeAfterSum s axis)

Sum a statically nonempty tensor axis.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.reduceSum_specFwd {s : Spec.Shape} {axis : ℕ} [Spec.Shape.HasNonemptyAxis axis s] [s.WellFormed] {α : Type} [TorchLean.Storage α] [Context α] (input : TorchLean.Tensor α s) : (reduceSum s axis).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = TorchLean.Tensor.reduceSum axis input ⋯

Pure evaluation of an axis sum.

def NN.GraphSpec.DAG.PrimOp.reduceMean (s : Spec.Shape) (axis : ℕ) [Spec.Shape.HasNonemptyAxis axis s] [s.WellFormed] : PrimOp [s] (TorchLean.Tensor.shapeAfterSum s axis)

Average a statically nonempty tensor axis.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.reduceMean_specFwd {s : Spec.Shape} {axis : ℕ} [Spec.Shape.HasNonemptyAxis axis s] [s.WellFormed] {α : Type} [TorchLean.Storage α] [Context α] (input : TorchLean.Tensor α s) : (reduceMean s axis).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = TorchLean.Tensor.reduceMean axis input ⋯

Pure evaluation of an axis mean.

def NN.GraphSpec.DAG.PrimOp.swapAdjacentAtDepth (s : Spec.Shape) (depth : ℕ) : PrimOp [s] (s.swapAdjacentAtDepth depth)

Swap adjacent tensor axes at a statically chosen depth.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.swapAdjacentAtDepth_specFwd {α : Type} [TorchLean.Storage α] [Context α] (s : Spec.Shape) (depth : ℕ) (input : TorchLean.Tensor α s) : (swapAdjacentAtDepth s depth).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = input.swapAdjacentAxes depth

Pure evaluation of an adjacent-axis swap.

def NN.GraphSpec.DAG.PrimOp.select (shape : Spec.Shape) (axis : ℕ) [Spec.Shape.AxisInBounds axis shape] (index : Fin (shape.axisSize axis)) : PrimOp [shape] (shape.eraseAxis axis)

Select one bounded coordinate along an arbitrary tensor axis.

def NN.GraphSpec.DAG.PrimOp.broadcast {source target : Spec.Shape} (proof : source.CanBroadcastTo target) : PrimOp [source] target

Expand singleton or missing dimensions according to a checked broadcasting derivation.

theorem NN.GraphSpec.DAG.PrimOp.applyAdjacentSwaps_range_eq_replaceAxis (shape : Spec.Shape) (axis extent : ℕ) (hAxis : axis < shape.rank) : (Spec.Shape.dim extent (shape.eraseAxis axis)).applyAdjacentSwaps (List.range axis) = shape.replaceAxis axis extent

Moving a replacement extent from the front to axis realizes replaceAxis.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.replaceAxis_axisSize (shape : Spec.Shape) (axis : ℕ) [hAxis : Spec.Shape.AxisInBounds axis shape] : shape.replaceAxis axis (shape.axisSize axis) = shape

Replacing an in-bounds axis by its existing extent leaves the shape unchanged.

def NN.GraphSpec.DAG.PrimOp.Internal.moveAxisToFront {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (shape : Spec.Shape) (axis extent : ℕ) [Spec.Shape.AxisInBounds axis shape] (input : Runtime.Autograd.Torch.Ref (shape.replaceAxis axis extent)) : m (Runtime.Autograd.Torch.Ref (Spec.Shape.dim extent (shape.eraseAxis axis)))

Move a chosen runtime axis to the front, retaining the unchanged axes in their order.

def NN.GraphSpec.DAG.PrimOp.Internal.moveAxisFromFront {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (shape : Spec.Shape) (axis extent : ℕ) [Spec.Shape.AxisInBounds axis shape] (input : Runtime.Autograd.Torch.Ref (Spec.Shape.dim extent (shape.eraseAxis axis))) : m (Runtime.Autograd.Torch.Ref (shape.replaceAxis axis extent))

Restore a front runtime axis to its original position.

def NN.GraphSpec.DAG.PrimOp.concatAxisSpec {α : Type} [TorchLean.Storage α] (shape : Spec.Shape) (axis left right : ℕ) [Spec.Shape.AxisInBounds axis shape] (a : TorchLean.Tensor α (shape.replaceAxis axis left)) (b : TorchLean.Tensor α (shape.replaceAxis axis right)) : TorchLean.Tensor α (shape.replaceAxis axis (left + right))

Concatenate tensors along an arbitrary statically valid axis.

def NN.GraphSpec.DAG.PrimOp.concatAxis (shape : Spec.Shape) (axis left right : ℕ) [Spec.Shape.AxisInBounds axis shape] : PrimOp [shape.replaceAxis axis left, shape.replaceAxis axis right] (shape.replaceAxis axis (left + right))

Concatenate along axis; shape supplies every unchanged axis extent.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.concatAxis_specFwd {shape : Spec.Shape} {axis left right : ℕ} [Spec.Shape.AxisInBounds axis shape] {α : Type} [TorchLean.Storage α] [Context α] (a : TorchLean.Tensor α (shape.replaceAxis axis left)) (b : TorchLean.Tensor α (shape.replaceAxis axis right)) : (concatAxis shape axis left right).specFwd (TorchLean.TensorPack.cons a (TorchLean.TensorPack.cons b TorchLean.TensorPack.nil)) = concatAxisSpec shape axis left right a b

The pure meaning of arbitrary-axis concatenation.

def NN.GraphSpec.DAG.PrimOp.sliceAxisRange (shape : Spec.Shape) (axis start length : ℕ) [Spec.Shape.AxisInBounds axis shape] (hRange : start + length ≤ shape.axisSize axis) : PrimOp [shape] (shape.replaceAxis axis length)

Select a checked contiguous range along an arbitrary statically valid axis.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.sliceAxisRange_specFwd {shape : Spec.Shape} {axis start length : ℕ} [Spec.Shape.AxisInBounds axis shape] (hRange : start + length ≤ shape.axisSize axis) {α : Type} [TorchLean.Storage α] [Context α] (input : TorchLean.Tensor α shape) : (sliceAxisRange shape axis start length hRange).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = TorchLean.Tensor.sliceAxisRangeSpec axis input start length hRange

The pure meaning of a checked arbitrary-axis contiguous slice.

def NN.GraphSpec.DAG.PrimOp.reshape (source target : Spec.Shape) (hSize : source.size = target.size) : PrimOp [source] target

Reshape a tensor while preserving its statically known number of entries.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.reshape_specFwd {source target : Spec.Shape} (hSize : source.size = target.size) {α : Type} [TorchLean.Storage α] [Context α] (input : TorchLean.Tensor α source) : (reshape source target hSize).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = input.reshapeSpec hSize

Pure evaluation of a shape-preserving view.