Executing Sequential Models

Execution operations for shape-checked sequential models and the typed graphs obtained by lowering them.

@[reducible, inline]

abbrev TorchLean.nn.TypedGraphModel (stateShapes : List Shape) (σ τ : Shape) (α : Type) [Storage α] : Type

A typed graph whose inputs are model state followed by one model input.

def TorchLean.nn.TypedGraphModel.forward {σ τ : Shape} {α : Type} [Storage α] {stateShapes : List Shape} (model : TypedGraphModel stateShapes σ τ α) (state : State α stateShapes) (input : Tensor α σ) : Tensor α τ

Evaluate a lowered model with explicit model state and one input tensor.

def TorchLean.nn.TypedGraphModel.jvp {σ τ : Shape} {α : Type} [Storage α] {stateShapes : List Shape} (model : TypedGraphModel stateShapes σ τ α) (state stateTangent : State α stateShapes) (input inputTangent : Tensor α σ) : Tensor α τ

Evaluate a Jacobian-vector product with separate tangents for the model state and input.

def TorchLean.nn.TypedGraphModel.vjp {σ τ : Shape} {α : Type} [Storage α] [Add α] [Zero α] {stateShapes : List Shape} (model : TypedGraphModel stateShapes σ τ α) (state : State α stateShapes) (input : Tensor α σ) (outputGradient : Tensor α τ) : State α stateShapes × Tensor α σ

Evaluate a vector-Jacobian product.

The returned pair is (stateGrad, inputGrad).