Typed GraphSpec DAG syntax

This module defines the shape-indexed variables, terms, argument lists, substitutions, and multi-result blocks used by the canonical GraphSpec DAG representation.

Primitives (arbitrary arity)

structure NN.GraphSpec.DAG.PrimOp (ins : List Spec.Shape) (τ : Spec.Shape) : Type 1

An n-ary primitive operation.

Compared to the sequential GraphSpec.Primitive, a PrimOp here is parameter-free: parameters are just ordinary inputs in the environment. This is what makes the DAG language flexible: a “layer” is expressed by let1-binding its parameters and then using them as inputs to ops.

Type indices:

• ins : List Shape is the ordered list of input tensor shapes the op expects.
• τ : Shape is the output tensor shape.

• name : String

  Debug name for error messages / inspection.

• specFwd {α : Type} [TorchLean.Storage α] [Context α] : TorchLean.TensorPack α ins → TorchLean.Tensor α τ

  Pure reference semantics (ins arguments packed as a typed list).

• program {α : Type} [TorchLean.Storage α] [Context α] : Runtime.Autograd.Model.Program α ins τ

  Executable TorchLean program with arguments of shapes ins.

Typed variables

inductive NN.GraphSpec.DAG.Var (Γ : List Spec.Shape) : Spec.Shape → Type

A de Bruijn variable whose result shape is part of its type.

Unlike a bare Fin Γ.length, Var Γ s records that the selected entry of Γ has shape s. Environment lookup therefore reduces structurally, without dependent casts through List.get. ofFin remains available to programmatic lowerings that discover positions dynamically.

• head {s : Spec.Shape} {Γ : List Spec.Shape} : Var (s :: Γ) s

  The first value in a nonempty environment.

• tail {Γ : List Spec.Shape} {s t : Spec.Shape} : Var Γ t → Var (s :: Γ) t

  A variable inherited from the tail of an environment.

def NN.GraphSpec.DAG.Var.ofFin {Γ : List Spec.Shape} (i : Fin Γ.length) : Var Γ (Γ.get i)

Convert a numeric environment position to a shape-indexed variable.

def NN.GraphSpec.DAG.Var.weakenRight {Γ : List Spec.Shape} {s t : Spec.Shape} : Var Γ s → Var (Γ ++ [t]) s

Preserve a variable when one value is appended to its environment.

def NN.GraphSpec.DAG.Var.inLeft {Γ : List Spec.Shape} (right : List Spec.Shape) {s : Spec.Shape} : Var Γ s → Var (Γ ++ right) s

Embed a variable from the left side of an appended environment.

def NN.GraphSpec.DAG.Var.inRight (left : List Spec.Shape) {Γ : List Spec.Shape} {s : Spec.Shape} : Var Γ s → Var (left ++ Γ) s

Embed a variable from the right side of an appended environment.

def NN.GraphSpec.DAG.Var.last (Γ : List Spec.Shape) {s : Spec.Shape} : Var (Γ ++ [s]) s

The final variable in an environment extended by one value.

def NN.GraphSpec.DAG.Var.liftRight {Γ Δ : List Spec.Shape} {t : Spec.Shape} (ρ : {s : Spec.Shape} → Var Γ s → Var Δ s) {s : Spec.Shape} (v : Var (Γ ++ [t]) s) : Var (Δ ++ [t]) s

Extend a shape-preserving variable renaming across one value appended to both environments.

DAG terms + arguments (mutual)

inductive NN.GraphSpec.DAG.Term (Γ : List Spec.Shape) : Spec.Shape → Type 1

A well-typed DAG term.

Read this as: “under environment Γ, this term produces a tensor of shape τ”.

• var {Γ : List Spec.Shape} {s : Spec.Shape} (i : Var Γ s) : Term Γ s

  Variable read (shape-indexed de Bruijn position in the environment).

• cast {Γ : List Spec.Shape} {σ τ : Spec.Shape} : Term Γ σ → σ = τ → Term Γ τ

  Cast a term’s output shape along a propositional equality.

  This is an internal hygiene tool: when we build terms programmatically (e.g. by lowering a higher-level syntax into DAG form), we often end up with goals like “Γ.get i = τ” that are true but not definitional.

  Using a cast node keeps the term in constructor form (so evaluators/compilers can still pattern match), and pushes the non-definitional equality into the semantics where it can be handled by cases h.

• castEnv {Γ Γ' : List Spec.Shape} {τ : Spec.Shape} : Term Γ τ → Γ = Γ' → Term Γ' τ

  Cast a term’s environment along a propositional equality.

  This is useful when normalizing list-association/parenthesization choices in Γ without changing meaning.

• op {Γ ins : List Spec.Shape} {τ : Spec.Shape} : PrimOp ins τ → Args Γ ins → Term Γ τ

  Apply an n-ary primitive op to n arguments.

• let1 {Γ : List Spec.Shape} {σ τ : Spec.Shape} : Term Γ σ → Term (Γ ++ [σ]) τ → Term Γ τ

  Let-bind a single intermediate value, extending the environment.

inductive NN.GraphSpec.DAG.Args (Γ : List Spec.Shape) : List Spec.Shape → Type 1

A typed list of argument terms.

Args Γ [s₁, …, sₙ] is a tuple of n terms, each well-typed under the same environment Γ, with corresponding shapes s₁, …, sₙ.

• nil {Γ : List Spec.Shape} : Args Γ []
• cons {Γ : List Spec.Shape} {s : Spec.Shape} {ss : List Spec.Shape} : Term Γ s → Args Γ ss → Args Γ (s :: ss)

def NN.GraphSpec.DAG.Args.rename {Γ Δ ss : List Spec.Shape} (ρ : {t : Spec.Shape} → Var Γ t → Var Δ t) : Args Γ ss → Args Δ ss

Rename every free variable in a typed operation-argument list.

def NN.GraphSpec.DAG.Term.rename {Γ Δ : List Spec.Shape} {s : Spec.Shape} (ρ : {t : Spec.Shape} → Var Γ t → Var Δ t) : Term Γ s → Term Δ s

Rename every free variable of a term while preserving its tensor shape.

def NN.GraphSpec.DAG.Term.weakenLeft {Γ : List Spec.Shape} {s t : Spec.Shape} : Term Γ s → Term (t :: Γ) s

Preserve a term when an unrelated value is prepended to its environment.

def NN.GraphSpec.DAG.Args.weakenLeft {Γ ss : List Spec.Shape} {t : Spec.Shape} : Args Γ ss → Args (t :: Γ) ss

Preserve typed operation arguments when a value is prepended to their environment.

def NN.GraphSpec.DAG.Term.weakenRight {Γ : List Spec.Shape} {s t : Spec.Shape} (term : Term Γ s) : Term (Γ ++ [t]) s

Preserve a term when an unrelated value is appended to its environment.

def NN.GraphSpec.DAG.Term.weakenAppend {Γ : List Spec.Shape} (right : List Spec.Shape) {s : Spec.Shape} (term : Term Γ s) : Term (Γ ++ right) s

Preserve a term when an arbitrary typed environment is appended.

def NN.GraphSpec.DAG.Args.get {Γ Δ : List Spec.Shape} {s : Spec.Shape} : Args Δ Γ → Var Γ s → Term Δ s

Select the term assigned to a shape-indexed variable.

def NN.GraphSpec.DAG.Args.append {Γ left right : List Spec.Shape} : Args Γ left → Args Γ right → Args Γ (left ++ right)

Concatenate two typed argument lists living in the same graph environment.

def NN.GraphSpec.DAG.Args.splitAppend {Γ left right : List Spec.Shape} : Args Γ (left ++ right) → Args Γ left × Args Γ right

Split arguments at a type-level list boundary.

This is the argument-list counterpart of TensorPack.splitAppend. It is useful when a model owns a concatenated parameter ABI but its implementation is assembled recursively from smaller models: each component receives exactly the terms belonging to its part of the ABI, with every tensor shape retained by the type checker.

@[simp]

theorem NN.GraphSpec.DAG.Args.splitAppend_append {Γ left right : List Spec.Shape} (leftArgs : Args Γ left) (rightArgs : Args Γ right) : (leftArgs.append rightArgs).splitAppend = (leftArgs, rightArgs)

Splitting arguments immediately after concatenating them recovers both original lists.

theorem NN.GraphSpec.DAG.Args.append_splitAppend {Γ left right : List Spec.Shape} (args : Args Γ (left ++ right)) : args.splitAppend.1.append args.splitAppend.2 = args

Concatenating both parts of a split recovers the original typed argument list.

def NN.GraphSpec.DAG.Args.vars (Γ : List Spec.Shape) : Args Γ Γ

View every entry of a typed environment as a term in that same environment.

The result preserves the order and shape indices of Γ. Large graph definitions can therefore pattern-match once on vars Γ instead of selecting every input by a numeric index and separately proving that the selected position has the expected shape.

@[simp]

theorem NN.GraphSpec.DAG.Args.get_rename {Γ Δ ss : List Spec.Shape} {s : Spec.Shape} (ρ : {t : Spec.Shape} → Var Γ t → Var Δ t) (args : Args Γ ss) (position : Var ss s) : (rename (fun {t : Spec.Shape} => ρ) args).get position = Term.rename (fun {t : Spec.Shape} => ρ) (args.get position)

Selecting from renamed arguments renames the selected term.

@[simp]

theorem NN.GraphSpec.DAG.Args.get_weakenLeft {Γ ss : List Spec.Shape} {s t : Spec.Shape} (args : Args Γ ss) (position : Var ss s) : args.weakenLeft.get position = (args.get position).weakenLeft

Lookup commutes with embedding an argument list below one new graph variable.

@[simp]

theorem NN.GraphSpec.DAG.Args.get_vars {Γ : List Spec.Shape} {s : Spec.Shape} (position : Var Γ s) : (vars Γ).get position = Term.var position

Selecting a variable from the complete environment returns that variable as a term.

Typed substitution

Substitution is the operation used to inline one graph into another. A substitution assigns a well-typed term in Δ to every variable in Γ; applying it replaces the free variables of a term without changing the term's result shape. The let-binding case lifts the assignment across the new value appended by the binder.

@[reducible, inline]

abbrev NN.GraphSpec.DAG.Substitution (Γ Δ : List Spec.Shape) : Type 1

A shape-preserving assignment of terms to every variable in an environment.

@[irreducible]

def NN.GraphSpec.DAG.Substitution.liftRight {Γ Δ : List Spec.Shape} {t : Spec.Shape} (σ : Substitution Γ Δ) : Substitution (Γ ++ [t]) (Δ ++ [t])

Extend a substitution across one value appended to both environments.

def NN.GraphSpec.DAG.Args.substitute {Γ Δ shapes : List Spec.Shape} (σ : Substitution Γ Δ) : Args Γ shapes → Args Δ shapes

Replace every free variable in a typed operation-argument list.

def NN.GraphSpec.DAG.Term.substitute {Γ Δ : List Spec.Shape} {s : Spec.Shape} (σ : Substitution Γ Δ) : Term Γ s → Term Δ s

Replace every free variable in a term by its assigned term.

def NN.GraphSpec.DAG.Term.instantiate {Γ Δ : List Spec.Shape} {s : Spec.Shape} (arguments : Args Δ Γ) (term : Term Γ s) : Term Δ s

Inline a term by supplying one typed argument term for each free variable.

inductive NN.GraphSpec.DAG.Block (Γ : List Spec.Shape) : List Spec.Shape → Type 1

A multi-result DAG block with shared A-normal-form bindings.

Block.let1 computes an intermediate once and makes it available to every eventual result. This is the representation needed by recurrent steps that return both a new state and an output derived from that state.

• ret {Γ outs : List Spec.Shape} : Args Γ outs → Block Γ outs

  Return one term for each output shape.

• let1 {Γ outs : List Spec.Shape} {σ : Spec.Shape} : Term Γ σ → Block (Γ ++ [σ]) outs → Block Γ outs

  Compute one shared intermediate and append it to the environment.

def NN.GraphSpec.DAG.Block.substitute {Γ Δ outputs : List Spec.Shape} (σ : Substitution Γ Δ) : Block Γ outputs → Block Δ outputs

Replace every free variable in a multi-result block.

def NN.GraphSpec.DAG.Block.instantiate {Γ Δ outputs : List Spec.Shape} (arguments : Args Δ Γ) (block : Block Γ outputs) : Block Δ outputs

Inline a multi-result block by supplying all of its free variables.

@[irreducible]

def NN.GraphSpec.DAG.Block.andThenWithRenaming {Γ Δ middle outputs : List Spec.Shape} (ρ : {s : Spec.Shape} → Var Γ s → Var Δ s) (first : Block Δ middle) (second : Block (Γ ++ middle) outputs) : Block Δ outputs

Compose two multi-output blocks.

The second block sees the original environment followed by every result of the first block. Any let1 bindings inside the first block remain shared. This is the typed DAG analogue of binding a tuple-valued computation and is the basic operation needed to compose recurrent cells, residual branches, and encoder-decoder stages.

def NN.GraphSpec.DAG.Block.andThen {Γ middle outputs : List Spec.Shape} (first : Block Γ middle) (second : Block (Γ ++ middle) outputs) : Block Γ outputs

Feed every output of first to second, preserving the original environment.