Canonical GraphSpec DAG

This file defines the canonical general GraphSpec representation: typed SSA/DAG terms.

NN.GraphSpec.Core gives a sequential authoring language (Graph with >>>) for chain-like architectures. That syntax lowers into this DAG language. Many modern architectures are not purely sequential:

• skip connections (ResNets),
• shared subcomputations (reusing the same feature tensor multiple times),
• multi-input ops (e.g. add, concatenation, attention blocks, …).

GraphSpec.DAG is the general version: a small SSA/A-normal-form term language whose terms denote DAG-shaped computation graphs.

Core idea: typed environments

We track an environment Γ : List Shape at the type level. A term Term Γ τ means:

• “if you provide values for every shape in Γ (in order),”
• “this term computes a tensor of shape τ.”

This is similar in spirit to a simply-typed lambda calculus with de Bruijn variables, except:

• there are no lambdas (only let-bindings), so the graph is acyclic by construction,
• primitive ops have arbitrary arity ins : List Shape,
• and Args Γ ins is a typed list of input subterms for an op.

The constructors are:

• var : read a variable from the environment (Fin Γ.length index),
• cast / castEnv : explicit casts to handle non-definitional equalities in Shape / Γ,
• op : apply an n-ary primitive to n arguments,
• let1 : bind an intermediate result and extend the environment (Γ ++ [σ]).

Mathematical semantics (intended)

For a fixed scalar type α with [Context α], we interpret an environment as a typed list of tensors TList α Γ. Then:

• Term.eval : TList α Γ → Term Γ τ → Tensor α τ is the pure reference semantics.
• Term.compile produces a backend-generic TorchLean program that computes the same graph, but in the executable (monadic, reference-based) runtime world.

Small example (residual add)

The residual pattern “main path + skip path” is the canonical example that needs DAG syntax:

y   = linear(W,b,x)
out = relu(y + x)

Here x is used twice, so a pure chain representation would need duplication. With let1, sharing is explicit: compute y once, then reuse it.

See NN/GraphSpec/Models/ResidualLinear.lean for a complete, well-typed instance.

References / citations

Conceptual background (stable, standard references):

• SSA form: Cytron et al. (1991), “Efficiently Computing Static Single Assignment Form…”.
• A-normal form / let-normal form (used to enforce DAG structure).
• Automatic differentiation survey: Baydin et al. (2018), “Automatic Differentiation in Machine Learning: a Survey”.
• Residual networks: He et al. (2016), “Deep Residual Learning for Image Recognition”.

Primitives (arbitrary arity)

structure NN.GraphSpec.DAG.PrimOp (ins : List Shape) (τ : 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} [Context α] : Runtime.Autograd.Torch.TList α ins → Tensor.Tensor α τ

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

• torchProgram {α : Type} [Context α] [DecidableEq Shape] : Runtime.Autograd.TorchLean.Program α ins τ

  Executable TorchLean program with arguments of shapes ins.

DAG terms + arguments (mutual)

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

A well-typed DAG term.

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

• var {Γ : List Shape} (i : Fin Γ.length) : Term Γ (Γ.get i)

  Variable read (de Bruijn index into the environment).

• cast {Γ : List Shape} {σ τ : 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 Shape} {τ : 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 Shape} {τ : Shape} : PrimOp ins τ → Args Γ ins → Term Γ τ

  Apply an n-ary primitive op to n arguments.

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

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

inductive NN.GraphSpec.DAG.Args (Γ : List Shape) : List 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 Shape} : Args Γ []
• cons {Γ : List Shape} {s : Shape} {ss : List Shape} : Term Γ s → Args Γ ss → Args Γ (s :: ss)

def NN.GraphSpec.DAG.Env.tget {α : Type} {ss : List Shape} : Runtime.Autograd.Torch.TList α ss → (i : Fin ss.length) → Tensor.Tensor α (ss.get i)

Typed environment lookup for pure tensors.

This is the underlying “variable semantics” for Term.eval.

def NN.GraphSpec.DAG.Env.rget {Ref : Shape → Type} {ss : List Shape} : Runtime.Autograd.Torch.RefList Ref ss → (i : Fin ss.length) → Ref (ss.get i)

Typed environment lookup for backend references.

This is the underlying “variable semantics” for Term.compile.

Spec interpreter

def NN.GraphSpec.DAG.Term.evalArgs {Γ ins : List Shape} {α : Type} [Context α] (env : Runtime.Autograd.Torch.TList α Γ) : Args Γ ins → Runtime.Autograd.Torch.TList α ins

Evaluate a typed argument list by evaluating each component term under the same environment.

def NN.GraphSpec.DAG.Term.eval {Γ : List Shape} {τ : Shape} {α : Type} [Context α] (env : Runtime.Autograd.Torch.TList α Γ) : Term Γ τ → Tensor.Tensor α τ

Pure evaluation of a DAG term.

This is the “math-first” semantics: we interpret a term as a pure function on tensors. No monads, no mutation, no autograd tape — just the Spec definitions of primitives.

The key runtime discipline is the environment discipline:

• var reads from env,
• op evaluates its arguments and feeds them to the primitive’s specFwd,
• let1 evaluates the bound term once and extends env for the body.

TorchLean compiler

@[reducible, inline]

abbrev NN.GraphSpec.DAG.Term.RefT (m : Type → Type) (α : Type) [Context α] [DecidableEq Shape] [Runtime.Autograd.Torch.Ops m α] (s : Shape) : Type

RefT is the backend’s reference type for tensors of a given shape.

In the executable runtime, primitives operate on references (allocated tensors) inside a monad. This matches how typical deep-learning runtimes model device placement, mutation, and autograd.

def NN.GraphSpec.DAG.Term.compileArgs {Γ ins : List Shape} {α : Type} [Context α] [DecidableEq Shape] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (env : Runtime.Autograd.Torch.RefList (RefT m α) Γ) : Args Γ ins → m (Runtime.Autograd.Torch.RefList (RefT m α) ins)

Compile a typed argument list by compiling each component term under the same environment.

def NN.GraphSpec.DAG.Term.compile {Γ : List Shape} {τ : Shape} {α : Type} [Context α] [DecidableEq Shape] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (env : Runtime.Autograd.Torch.RefList (RefT m α) Γ) : Term Γ τ → m (RefT m α τ)

Compile a typed Term Γ τ into the backend monad m, producing a reference to a tensor of shape τ.

This is the "executable" counterpart of Term.eval: instead of returning a pure Tensor, we emit backend ops (Runtime.Autograd.Torch.Ops) that allocate tensors and apply primitives.

Model wrapper

structure NN.GraphSpec.DAG.Model (ps ins : List Shape) (τ : Shape) : Type 1

A small “model” wrapper around DAG terms.

This mirrors the sequential Graph surface:

• ps are parameter tensor shapes (tracked at the type level),
• ins are the shapes of non-parameter inputs (e.g. data tensors),
• τ is the output shape.

The model body is a Term (ps ++ ins) τ, i.e. it expects an environment that starts with parameters and then contains the actual inputs.

• initParams : Runtime.Autograd.Torch.TList Float ps

  init Params.

• body : Term (ps ++ ins) τ

  body.

def NN.GraphSpec.DAG.Model.specFwd {ps ins : List Shape} {τ : Shape} (m : Model ps ins τ) {α : Type} [Context α] (params : Runtime.Autograd.Torch.TList α ps) (xs : Runtime.Autograd.Torch.TList α ins) : Tensor.Tensor α τ

Pure forward semantics of a DAG model.

We build the full environment Γ = ps ++ ins by appending the parameter list and the input list, then evaluate the body using Term.eval.

def NN.GraphSpec.DAG.Model.torchProgram {ps ins : List Shape} {τ : Shape} (m : Model ps ins τ) {α : Type} [Context α] [DecidableEq Shape] : Runtime.Autograd.TorchLean.Program α (ps ++ ins) τ

Compile a DAG model to a backend-generic TorchLean program.

The resulting program expects arguments in the order ps ++ ins (parameters first, then inputs), matching the environment discipline used by specFwd.