Model State IO

This module saves and restores the complete shape-indexed state of a TorchLean runtime module. State includes trainable parameters and persistent buffers; optimizer state has a separate format.

Two formats are provided for runtime state packs.

Exact Scalar Format

Checkpoint.Encoding selects an exact scalar payload and a versioned format tag. Native and configured binary values use storage words; complex values store both component words. The Float instance retains the existing binary64 format, so old checkpoints remain readable. The streamed float32 format further down remains a separate device-oriented representation.

This is:

• exact for exported scalar encodings (including subnormals and configured NaN payloads),
• stable across locales, and
• easy to validate (length = Spec.Shape.size).

The file layout is:

{
  "format": "torchlean_state_bits_v1",
  "state": [
    { "shape": [d1, d2, ...], "values": [u64bits, u64bits, ...] },
    ...
  ]
}

Runtime Float32 Format

Native Float32 modules and device-backed modules use a versioned binary stream. Each state tensor records its rank, dimensions, and element count before its little-endian binary32 payload. Loading checks all metadata against the expected shape-indexed state and rejects unsupported versions, truncated payloads, and trailing data. The stream is written one tensor at a time, so saving a large CUDA model does not construct a second host-side copy of the whole checkpoint.

def Runtime.Autograd.Model.StateIO.float32StreamFormat : Torch.Internal.CheckpointIO.Format

Versioned header for streamed runtime float32 state checkpoints.

def Runtime.Autograd.Model.StateIO.jsonNat (n : ℕ) : Lean.Json

Encode a natural number as a JSON number.

def Runtime.Autograd.Model.StateIO.tensorToJsonBits {α : Type} [TorchLean.Storage α] [TorchLean.Checkpoint.Encoding α] (s : Spec.Shape) (t : TorchLean.Tensor α s) : Lean.Json

Encode one tensor as shape metadata plus exact scalar payloads.

def Runtime.Autograd.Model.StateIO.tensorFromJsonBits {α : Type} [TorchLean.Storage α] [TorchLean.Checkpoint.Encoding α] (tag : String) (s : Spec.Shape) (j : Lean.Json) : Except String (TorchLean.Tensor α s)

Decode one shape-checked tensor from the exact scalar state format.

def Runtime.Autograd.Model.StateIO.stateToJsonBits {α : Type} [TorchLean.Storage α] [TorchLean.Checkpoint.Encoding α] {ss : List Spec.Shape} : TorchLean.TensorPack α ss → Lean.Json

Encode shape-indexed model state as the JSON array stored under state.

def Runtime.Autograd.Model.StateIO.stateFromJsonBitsArray {α : Type} [TorchLean.Storage α] [TorchLean.Checkpoint.Encoding α] (tag : String) (xs : Array Lean.Json) {ss : List Spec.Shape} (offset : ℕ) : Except String (TorchLean.TensorPack α ss)

Decode an expected state layout from a tensor array, starting at offset.

def Runtime.Autograd.Model.StateIO.stateFromJsonBits {α : Type} [TorchLean.Storage α] [TorchLean.Checkpoint.Encoding α] (tag : String) {ss : List Spec.Shape} (json : Lean.Json) : Except String (TorchLean.TensorPack α ss)

Decode the state JSON array into the expected shape-indexed state pack.

def Runtime.Autograd.Model.StateIO.writeStateBits {α : Type} [TorchLean.Storage α] [TorchLean.Checkpoint.Encoding α] (path : System.FilePath) {ss : List Spec.Shape} (state : TorchLean.TensorPack α ss) (pretty : Bool := true) : IO Unit

Write model state using the scalar type's exact encoding and versioned format tag.

def Runtime.Autograd.Model.StateIO.readStateBits {α : Type} [TorchLean.Storage α] [TorchLean.Checkpoint.Encoding α] (path : System.FilePath) {ss : List Spec.Shape} : IO (Except String (TorchLean.TensorPack α ss))

Read model state, rejecting incompatible scalar formats, shapes, and payload lengths.

Streaming float32 module checkpoints

def Runtime.Autograd.Model.StateIO.pushUInt32LE (out : ByteArray) (bits : UInt32) : ByteArray

Append one binary32 bit pattern in little-endian byte order.

def Runtime.Autograd.Model.StateIO.uint32FromLE (bytes : ByteArray) (offset : ℕ) : UInt32

Decode one little-endian 32-bit word at offset.

def Runtime.Autograd.Model.StateIO.tensorToFloat32Bytes {α : Type} [TorchLean.Storage α] (encode : α → Float32) {shape : Spec.Shape} (tensor : TorchLean.Tensor α shape) : ByteArray

Encode a tensor as exact little-endian binary32 values.

def Runtime.Autograd.Model.StateIO.tensorFromFloat32Bytes {α : Type} [TorchLean.Storage α] (decode : Float32 → α) (shape : Spec.Shape) (bytes : ByteArray) : Except String (TorchLean.Tensor α shape)

Decode an exact binary32 payload into a shape-checked tensor.

def Runtime.Autograd.Model.StateIO.writeShape (handle : IO.FS.Handle) (shape : Spec.Shape) : IO Unit

Write one expected tensor shape to a streaming checkpoint.

def Runtime.Autograd.Model.StateIO.readShape (handle : IO.FS.Handle) (expected : Spec.Shape) : IO Unit

Read and validate one tensor shape from a streaming checkpoint.

def Runtime.Autograd.Model.StateIO.tensorFloat32Bytes {α : Type} [TorchLean.Storage α] (encode : α → Float32) {shape : Spec.Shape} (tensorRef : Torch.Param α shape) : IO ByteArray

Obtain one state tensor as raw float32 bytes without materializing the whole state pack.

def Runtime.Autograd.Model.StateIO.writeStateFloat32 {α : Type} [TorchLean.Storage α] (encode : α → Float32) (handle : IO.FS.Handle) {shapes : List Spec.Shape} : Torch.ParamList α shapes → IO Unit

Stream shape-indexed runtime state to an open checkpoint handle.

def Runtime.Autograd.Model.StateIO.writeModuleStateFloat32 {α : Type} [TorchLean.Storage α] (encode : α → Float32) (path : System.FilePath) {shapes : List Spec.Shape} (state : Torch.ParamList α shapes) : IO Unit

Write runtime model state as a streamed float32 checkpoint.

Only one tensor payload is resident on the host at a time. This is the appropriate format for large CUDA models; it records the exact values used by the float32 runtime instead of expanding them into one in-memory JSON tree.

def Runtime.Autograd.Model.StateIO.isModuleStateFloat32 (path : System.FilePath) : IO Bool

Check whether a file begins with the streaming float32 checkpoint header.

def Runtime.Autograd.Model.StateIO.readTensorFloat32Into {α : Type} [TorchLean.Storage α] (decode : Float32 → α) (useCuda : Bool) (handle : IO.FS.Handle) {shape : Spec.Shape} (tensorRef : Torch.Param α shape) : IO Unit

Read one float32 payload directly into an existing runtime state tensor.

def Runtime.Autograd.Model.StateIO.readStateFloat32Into {α : Type} [TorchLean.Storage α] (decode : Float32 → α) (useCuda : Bool) (handle : IO.FS.Handle) {shapes : List Spec.Shape} : Torch.ParamList α shapes → IO Unit

Stream a checkpoint into shape-indexed runtime state.

def Runtime.Autograd.Model.StateIO.validateStateFloat32 (handle : IO.FS.Handle) (shapes : List Spec.Shape) : IO Unit

Validate every tensor record and payload without changing the destination parameters.

def Runtime.Autograd.Model.StateIO.readAndCheckFloat32Header (handle : IO.FS.Handle) (expectedTensorCount : ℕ) : IO Unit

Read and validate the common header of a streamed float32 checkpoint.

def Runtime.Autograd.Model.StateIO.checkFloat32EndOfFile (handle : IO.FS.Handle) : IO Unit

Reject extra bytes after the last expected tensor payload.

def Runtime.Autograd.Model.StateIO.readModuleStateFloat32Into {α : Type} [TorchLean.Storage α] (decode : Float32 → α) (path : System.FilePath) (useCuda : Bool) {shapes : List Spec.Shape} (state : Torch.ParamList α shapes) : IO Unit

Load a streamed float32 checkpoint into existing runtime state.