Core DAG Primitives

Constants, elementwise arithmetic, and standard sequential operations exposed as typed DAG nodes.

Basic DAG primitives

def NN.GraphSpec.DAG.PrimOp.zero (s : Spec.Shape) : PrimOp [] s

Produce the all-zero tensor of a statically known shape.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.zero_specFwd (s : Spec.Shape) {α : Type} [TorchLean.Storage α] [Context α] : (zero s).specFwd TorchLean.TensorPack.nil = TorchLean.Tensor.full s 0

The zero DAG node denotes the all-zero tensor of its declared shape.

def NN.GraphSpec.DAG.PrimOp.one (s : Spec.Shape) : PrimOp [] s

Produce the all-one tensor of a statically known shape.

def NN.GraphSpec.DAG.PrimOp.linear (inDim outDim : ℕ) : PrimOp [[outDim, inDim], [outDim], [inDim]] [outDim]

Dense linear layer in DAG form.

Inputs are ordered as [W, b, x]:

• W : Tensor α [outDim, inDim],
• b : Tensor α [outDim],
• x has shape [inDim].

The output has shape [outDim]. This is the DAG embedding of Primitive.linear, so the DAG and sequential authoring surfaces share the same Spec semantics and TorchLean lowering path.

def NN.GraphSpec.DAG.PrimOp.flatten (s : Spec.Shape) : PrimOp [s] [s.size]

Flatten a tensor to a rank-one tensor in DAG form.

Input: [x : TorchLean.Tensor s]. Output: Tensor α [Spec.Shape.size s].

This is the DAG embedding of Primitive.flatten, so it has exactly the same row-major view semantics as the sequential primitive.

Spatial and residual DAG primitives

def NN.GraphSpec.DAG.PrimOp.relu (s : Spec.Shape) : PrimOp [s] s

ReLU activation in DAG form.

Input: [x : s], output: s.

Semantics: elementwise $\max(x,0)$. This is parameter-free and derived from Primitive.relu.

Reference: Nair and Hinton (2010), "Rectified Linear Units Improve Restricted Boltzmann Machines".

def NN.GraphSpec.DAG.PrimOp.add (s : Spec.Shape) : PrimOp [s, s] s

Add two tensors of the same shape.

Input shapes: [s, s], output shape: s.

This is the primitive used for residual/skip connections: $\mathrm{out}=\operatorname{main}(x)+x$. It is defined directly because the sequential surface is unary, while residual addition is genuinely multi-input.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.add_specFwd {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (left right : TorchLean.Tensor α s) : (add s).specFwd (TorchLean.TensorPack.cons left (TorchLean.TensorPack.cons right TorchLean.TensorPack.nil)) = left.addSpec right

The pure meaning of the DAG addition node is pointwise tensor addition.

def NN.GraphSpec.DAG.PrimOp.sub (s : Spec.Shape) : PrimOp [s, s] s

Subtract two tensors of the same shape.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.sub_specFwd {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (left right : TorchLean.Tensor α s) : (sub s).specFwd (TorchLean.TensorPack.cons left (TorchLean.TensorPack.cons right TorchLean.TensorPack.nil)) = left.subSpec right

The pure meaning of the DAG subtraction node is pointwise tensor subtraction.

def NN.GraphSpec.DAG.PrimOp.mul (s : Spec.Shape) : PrimOp [s, s] s

Multiply two tensors coordinatewise.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.mul_specFwd {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (left right : TorchLean.Tensor α s) : (mul s).specFwd (TorchLean.TensorPack.cons left (TorchLean.TensorPack.cons right TorchLean.TensorPack.nil)) = left.mulSpec right

The pure meaning of the DAG multiplication node is pointwise tensor multiplication.

def NN.GraphSpec.DAG.PrimOp.conv {d : ℕ} (inC outC : ℕ) (kernel stride padding spatial : TorchLean.Tensor ℕ [d]) : PrimOp [Spec.Shape.ofList (outC :: inC :: kernel.to (List ℕ)), [outC], Spec.Shape.ofList (inC :: spatial.to (List ℕ))] (Spec.Shape.ofList (outC :: (Spec.convOutSpatial spatial kernel stride padding).to (List ℕ)))

Arbitrary-rank convolution in DAG form, with inputs ordered as [kernel, bias, x].

def NN.GraphSpec.DAG.PrimOp.maxPool {d : ℕ} (channels : ℕ) (kernel stride padding spatial : TorchLean.Tensor ℕ [d]) {hKernel : ∀ (i : Fin d), kernel.getScalar i ≠ 0} {hStride : ∀ (i : Fin d), stride.getScalar i ≠ 0} : PrimOp [Spec.Shape.ofList (channels :: spatial.to (List ℕ))] (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial kernel stride padding).to (List ℕ)))

Arbitrary-rank max pooling in DAG form.

def NN.GraphSpec.DAG.PrimOp.batchNorm (channels : ℕ) (spatial : Spec.Shape) (hWellFormed : (Spec.Shape.dim channels spatial).wellFormed) : PrimOp [[channels], [channels], Spec.Shape.dim channels spatial] (Spec.Shape.dim channels spatial)

Batch normalization over an arbitrary spatial shape in DAG form.