GraphSpec DAG model wrappers

This module packages typed terms and blocks as parameterized single- and multi-output models, with pure forward semantics, inlining theorems, and executable program conversion.

Model wrapper

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

A small “model” wrapper around DAG terms.

This mirrors the sequential Chain 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 : TorchLean.TensorPack Float ps

  Initial parameter values, one tensor per shape in ps. Shipping the initialization with the model means a Model is runnable on its own, with no separate setup step.

• body : Term (ps ++ ins) τ

  The computation itself, as a term over the environment ps ++ ins: parameters first, then data inputs. That fixed ordering is what lets initParams be typed by ps alone.

def NN.GraphSpec.DAG.Model.inline {Γ ps ins : List Spec.Shape} {τ : Spec.Shape} (model : Model ps ins τ) (params : Args Γ ps) (inputs : Args Γ ins) : Term Γ τ

Inline a model body into a larger graph using explicit parameter and input terms.

The result is still an ordinary DAG term: no primitive boundary is introduced, and subsequent lowering, differentiation, or numerical analysis can inspect every operation of the model.

def NN.GraphSpec.DAG.Model.specFwd {ps ins : List Spec.Shape} {τ : Spec.Shape} (m : Model ps ins τ) {α : Type} [TorchLean.Storage α] [Context α] (params : TorchLean.TensorPack α ps) (xs : TorchLean.TensorPack α ins) : TorchLean.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.

theorem NN.GraphSpec.DAG.Model.eval_inline {Γ ps ins : List Spec.Shape} {τ : Spec.Shape} (model : Model ps ins τ) {α : Type} [TorchLean.Storage α] [Context α] (env : TorchLean.TensorPack α Γ) (params : Args Γ ps) (inputs : Args Γ ins) : Term.eval env (model.inline params inputs) = model.specFwd (Term.evalArgs env params) (Term.evalArgs env inputs)

Inlining a model into a larger DAG preserves the model's pure forward semantics.

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

Lower a DAG model to an execution-polymorphic TorchLean program.

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

Models with several outputs

structure NN.GraphSpec.DAG.MultiModel (ps ins outs : List Spec.Shape) : Type 1

A typed DAG model that returns several tensors.

Recurrent layers commonly return both an updated state and an observable output. Keeping these as a typed list preserves their individual shapes and avoids flattening unrelated tensors into an untyped buffer. The block body uses shared let1 bindings, so an updated state can be computed once and returned alongside values derived from it.

• initParams : TorchLean.TensorPack Float ps

  Default parameters, in the same ABI order used by body.

• body : Block (ps ++ ins) outs

  Shared computation ending in one well-typed term for every output shape.

def NN.GraphSpec.DAG.MultiModel.inline {Γ ps ins outs : List Spec.Shape} (model : MultiModel ps ins outs) (params : Args Γ ps) (inputs : Args Γ ins) : Block Γ outs

Inline a multi-output model into a larger graph.

Shared let bindings in the original block remain shared after substitution, so recurrent state updates are not duplicated when both the state and a derived output are returned.

def NN.GraphSpec.DAG.MultiModel.specFwd {ps ins outs : List Spec.Shape} (m : MultiModel ps ins outs) {α : Type} [TorchLean.Storage α] [Context α] (params : TorchLean.TensorPack α ps) (xs : TorchLean.TensorPack α ins) : TorchLean.TensorPack α outs

Pure reference semantics of a multi-output DAG model.

theorem NN.GraphSpec.DAG.MultiModel.eval_inline {Γ ps ins outs : List Spec.Shape} (model : MultiModel ps ins outs) {α : Type} [TorchLean.Storage α] [Context α] (env : TorchLean.TensorPack α Γ) (params : Args Γ ps) (inputs : Args Γ ins) : Block.eval env (model.inline params inputs) = model.specFwd (Term.evalArgs env params) (Term.evalArgs env inputs)

Inlining a multi-output model preserves every output and every shared intermediate in its pure reference semantics.

@[reducible, inline]

abbrev NN.GraphSpec.DAG.MultiModel.MultiOutputProgram (α : Type) [TorchLean.Storage α] [Context α] (ins outs : List Spec.Shape) : Type 1

An execution-polymorphic program returning several shape-indexed tensor references.

def NN.GraphSpec.DAG.MultiModel.toProgram {ps ins outs : List Spec.Shape} (m : MultiModel ps ins outs) {α : Type} [TorchLean.Storage α] [Context α] : MultiOutputProgram α (ps ++ ins) outs

Lower every result of a multi-output model for the selected TorchLean execution target.