Model State

nn.State α shapes is the public, shape-indexed container for model parameters and persistent buffers. Its recursive tensor-pack representation stays behind the runtime boundary; application code reads and replaces entries through named operations.

structure TorchLean.nn.State (α : Type) [Storage α] (shapes : List Shape) : Type

Model parameters and persistent buffers in their statically known order.

The element type is shared while every tensor keeps its own shape from shapes.

• tensors : TorchLean.TensorPack α shapes

def TorchLean.nn.State.Internal.fromTensorPack {α : Type} [Storage α] {shapes : List Shape} (tensors : TensorPack α shapes) : State α shapes

Wrap the runtime tensor-pack representation at an implementation boundary.

def TorchLean.nn.State.Internal.toTensorPack {α : Type} [Storage α] {shapes : List Shape} (state : State α shapes) : TensorPack α shapes

Reveal the runtime tensor-pack representation at an implementation boundary.

@[simp]

theorem TorchLean.nn.State.Internal.toTensorPack_fromTensorPack {α : Type} [Storage α] {shapes : List Shape} (tensors : TensorPack α shapes) : toTensorPack (fromTensorPack tensors) = tensors

Revealing a freshly wrapped tensor pack returns the original pack.

@[simp]

theorem TorchLean.nn.State.Internal.fromTensorPack_toTensorPack {α : Type} [Storage α] {shapes : List Shape} (state : State α shapes) : fromTensorPack (toTensorPack state) = state

Wrapping the representation of a state reconstructs that state.

def TorchLean.nn.State.Internal.replaceTensor {α : Type} [Storage α] {ss : List Shape} (tensors : TensorPack α ss) (index : Fin ss.length) (value : Tensor α (ss.get index)) : TensorPack α ss

Replace one tensor in the runtime representation.

theorem TorchLean.nn.State.ext {α : Type} [Storage α] {shapes : List Shape} {left right : State α shapes} (h : Internal.toTensorPack left = Internal.toTensorPack right) : left = right

Two states are equal when their tensor-pack representations are equal.

theorem TorchLean.nn.State.ext_iff {α : Type} [Storage α] {shapes : List Shape} {left right : State α shapes} : left = right ↔ Internal.toTensorPack left = Internal.toTensorPack right

def TorchLean.nn.State.empty {α : Type} [Storage α] : State α []

Empty model state.

def TorchLean.nn.State.full {α : Type} [Storage α] {shapes : List Shape} (value : α) : State α shapes

Construct state whose every tensor contains value.

def TorchLean.nn.State.zeros {α : Type} [Storage α] {shapes : List Shape} [Zero α] : State α shapes

Construct all-zero state.

def TorchLean.nn.State.get {α : Type} [Storage α] {shapes : List Shape} (state : State α shapes) (index : Fin shapes.length) : Tensor α (shapes.get index)

Read one state tensor; its result shape is determined by the index.

def TorchLean.nn.State.set {α : Type} [Storage α] {shapes : List Shape} (state : State α shapes) (index : Fin shapes.length) (value : Tensor α (shapes.get index)) : State α shapes

Replace one state tensor with another tensor of exactly the required shape.

def TorchLean.nn.State.map {α β : Type} [Storage α] [Storage β] {shapes : List Shape} (state : State α shapes) (f : {shape : Shape} → Tensor α shape → Tensor β shape) : State β shapes

Apply a shape-preserving conversion to every state tensor.

def TorchLean.nn.State.zipWith {α β γ : Type} [Storage α] [Storage β] [Storage γ] {shapes : List Shape} (first : State α shapes) (second : State β shapes) (f : {shape : Shape} → Tensor α shape → Tensor β shape → Tensor γ shape) : State γ shapes

Combine corresponding state tensors with a shape-preserving operation.

def TorchLean.nn.State.append {α : Type} [Storage α] {leftShapes rightShapes : List Shape} (first : State α leftShapes) (second : State α rightShapes) : State α (leftShapes ++ rightShapes)

Concatenate two model states while retaining the combined shape layout in the type.

def TorchLean.nn.State.push {α : Type} [Storage α] {shapes : List Shape} {shape : Shape} (state : State α shapes) (tensor : Tensor α shape) : State α (shapes ++ [shape])

Add one tensor at the end of a state, extending its statically known layout.

@[instance_reducible]

instance TorchLean.nn.State.instRepr {α : Type} [Storage α] {shapes : List Shape} [Repr α] : Repr (State α shapes)

structure TorchLean.nn.State.Partition (α : Type) [Storage α] (leftShapes rightShapes : List Shape) : Type

Named result of splitting state at a statically known shape-list boundary.

• left : State α leftShapes

  State before the split boundary.

• right : State α rightShapes

  State after the split boundary.

def TorchLean.nn.State.instReprPartition.repr {α✝ : Type} {inst✝ : Storage α✝} {leftShapes✝ rightShapes✝ : List Shape} [Repr α✝] : Partition α✝ leftShapes✝ rightShapes✝ → ℕ → Std.Format

@[instance_reducible]

instance TorchLean.nn.State.instReprPartition {α✝ : Type} {inst✝ : Storage α✝} {leftShapes✝ rightShapes✝ : List Shape} [Repr α✝] : Repr (Partition α✝ leftShapes✝ rightShapes✝)

def TorchLean.nn.State.split {α : Type} [Storage α] {leftShapes rightShapes : List Shape} (state : State α (leftShapes ++ rightShapes)) : Partition α leftShapes rightShapes

Split state at a statically known shape-list boundary.

def TorchLean.nn.State.cast {α : Type} [Storage α] {leftShapes rightShapes : List Shape} (state : State α leftShapes) (sameShapes : leftShapes = rightShapes) : State α rightShapes

Transport state along an equality between its statically known shape layouts.