Optimizer Checkpoint Schemas

An optimizer checkpoint is meaningful only for the parameter layout that produced it. This module stores that shared layout: parameter shapes in module order and the corresponding trainability flags. Optimizer-specific codecs validate this schema before installing backend state.

structure Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema : Type

Ordered parameter metadata shared by backend-owned optimizer checkpoints.

• shapes : Array Spec.Shape

  Parameter shapes in module order.

• requiresGrad : Array Bool

  Whether each corresponding parameter is trainable.

• representatives : Option (Array (Option ℕ))

  First trainable slot sharing each parameter's storage; frozen slots contain none.

  An absent map describes independent parameters, as in version 2 checkpoints. An explicit map retains the original slot numbers: [p, p, q] has representatives [0, 0, 2], so its moment entries have keys 0 and 2. The trainer constructs this map from storage identity.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.representative? (schema : ParameterSchema) (index : ℕ) : Option ℕ

Optimizer slot for a trainable parameter, or none for a frozen or invalid slot.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.isWellFormed (schema : ParameterSchema) : Bool

Check the ordered metadata and any explicit storage-sharing map.

A representative must be trainable, have the same shape, precede its aliases, and represent itself. These checks exclude cycles and chains, and keep frozen slots out of optimizer history.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.isRepresentative (schema : ParameterSchema) (index : ℕ) : Bool

Whether this slot owns a moment entry rather than referring to another slot's history.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.optimizerStateCount (schema : ParameterSchema) : ℕ

Number of distinct moment entries required by this layout.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.hasAliases (schema : ParameterSchema) : Bool

Whether any trainable slot shares an earlier slot's optimizer history.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.writeRepresentatives (format : CheckpointIO.Format) (handle : IO.FS.Handle) (schema : ParameterSchema) : IO Unit

Write the version 3 sharing map after the ordinary shape/trainability schema.

Each slot uses one unsigned 64-bit field: zero for a frozen parameter and representative + 1 otherwise. The preceding schema already fixes the number of fields and their parameter order.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.readAndCheckRepresentatives (format : CheckpointIO.Format) (handle : IO.FS.Handle) (expected : ParameterSchema) : IO Unit

Compare every stored representative with the layout derived from the destination parameters.

Matching just the number of moment entries would accept different sharing patterns such as [p, q, p] and [p, q, q]. Comparing all slots rejects that mismatch before moment allocation.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.write (format : CheckpointIO.Format) (handle : IO.FS.Handle) (schema : ParameterSchema) : IO Unit

Write the ordered parameter schema for an optimizer-specific checkpoint format.

def Runtime.Autograd.Torch.Internal.OptimizerCheckpoint.ParameterSchema.readAndCheck (format : CheckpointIO.Format) (handle : IO.FS.Handle) (expected : ParameterSchema) : IO Unit

Read a parameter schema and reject any difference from the expected module layout.