DAG Nonlinear and Attention Primitives

Elementwise nonlinearities, softmax, and multi-head self-attention for typed DAG models.

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

Coordinatewise sigmoid.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.sigmoid_specFwd {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (input : TorchLean.Tensor α s) : (sigmoid s).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = Activation.sigmoidSpec input

Pure evaluation of coordinatewise sigmoid.

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

Coordinatewise SiLU, also called swish: x * sigmoid x.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.silu_specFwd {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (input : TorchLean.Tensor α s) : (silu s).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = Activation.swishSpec input

Pure evaluation of coordinatewise SiLU.

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

Coordinatewise Gaussian error linear unit using TorchLean's tanh approximation.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.gelu_specFwd {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (input : TorchLean.Tensor α s) : (gelu s).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = Activation.geluSpec input

Pure evaluation of coordinatewise GELU.

def NN.GraphSpec.DAG.PrimOp.attention (leading : Spec.Shape) (n numHeads dModel headDim : ℕ) (hN : 0 < n) : PrimOp [[dModel, numHeads * headDim], [dModel, numHeads * headDim], [dModel, numHeads * headDim], [numHeads * headDim, dModel], leading.concat [n, dModel]] (leading.concat [n, dModel])

Multi-head self-attention over an arbitrary leading shape.

The pure semantics applies Spec.MultiHeadAttention.forward independently at every leading index. The executable path uses the prefix-polymorphic TorchLean operation and may flatten those axes for backend execution without changing the mathematical operation.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.attention_specFwd {leading : Spec.Shape} {n numHeads dModel headDim : ℕ} (hN : 0 < n) {α : Type} [TorchLean.Storage α] [Context α] (wq wk wv : TorchLean.Tensor α [dModel, numHeads * headDim]) (wo : TorchLean.Tensor α [numHeads * headDim, dModel]) (input : TorchLean.Tensor α (leading.concat [n, dModel])) : (attention leading n numHeads dModel headDim hN).specFwd (TorchLean.TensorPack.cons wq (TorchLean.TensorPack.cons wk (TorchLean.TensorPack.cons wv (TorchLean.TensorPack.cons wo (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil))))) = TorchLean.Tensor.mapLeading leading (fun (inputs : TorchLean.Tensor α [n, dModel]) => Spec.MultiHeadAttention.forward n ⋯ { queryWeight := wq, keyWeight := wk, valueWeight := wv, outputWeight := wo } inputs none) input

Pure evaluation of prefix-polymorphic multi-head attention.

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

Coordinatewise exponential.

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

Coordinatewise multiplicative inverse.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.inv_specFwd {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (input : TorchLean.Tensor α s) : (inv s).specFwd (TorchLean.TensorPack.cons input TorchLean.TensorPack.nil) = input.invSpec

The pure meaning of the DAG inverse node is coordinatewise reciprocal.

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

Coordinatewise hyperbolic tangent.

def NN.GraphSpec.DAG.PrimOp.softmax (s : Spec.Shape) (axis : ℕ) [Spec.Shape.AxisInBounds axis s] : PrimOp [s] s

Stable softmax along any in-bounds tensor dimension.