Eager Session State

The tapes and side tables shared by eager operations. Parameters own their CUDA mirrors; the session owns recorded intermediates. Handle generations separate consecutive forward passes.

structure Runtime.Autograd.Torch.Internal.ParameterStorage (α : Type) [TorchLean.Storage α] : Type

The mutable cells owned by a parameter, without reading or comparing tensor elements.

Two registrations share optimizer storage only when all three cells coincide. Equal values and equal names do not make separately allocated parameters aliases. Keeping this descriptor in the session leaves the public Param record and its construction interface unchanged.

• shape : Spec.Shape
• value : IO.Ref (TorchLean.Tensor α self.shape)
• cudaValue : IO.Ref (Option LibTorch.SomeBuffer)
• hostCurrent : IO.Ref Bool

def Runtime.Autograd.Torch.Internal.ParameterStorage.same {α : Type} [TorchLean.Storage α] (left right : ParameterStorage α) : IO Bool

Compare mutable-cell identity using Lean's safe stateful reference operation.

def Runtime.Autograd.Torch.Internal.ParameterStorage.canonicalIndices {α : Type} [TorchLean.Storage α] (storages : Array (Option (ParameterStorage α))) : IO (Array (Option ℕ))

Map each present storage to its first occurrence, retaining absent slots.

A balanced tree of full-width reference keys avoids pairwise searches among independent storages. Reference equality remains authoritative even if keys collide. Representatives retain their cells until the operation ends; neither keys nor references are cached between calls. Consequently a new parameter list or recording can change its alias layout without invalidating a cache.

inductive Runtime.Autograd.Torch.Internal.SessionTape (α : Type) [TorchLean.Storage α] : Type

The tape a session records on. Its values follow the session's device: shape-erased tensors on the CPU, device buffers on CUDA. A session has exactly one, made empty for options.device when the session opens and again at every reset.

• cpu {α : Type} [TorchLean.Storage α] (tape : Tape (Spec.SomeTensor α)) : SessionTape α

  The tape of a CPU session.

• cuda {α : Type} [TorchLean.Storage α] (tape : Tape LibTorch.SomeBuffer) : SessionTape α

  The tape of a CUDA session.

def Runtime.Autograd.Torch.Internal.SessionTape.empty {α : Type} [TorchLean.Storage α] (device : NN.Backend.Device) : SessionTape α

An empty tape for a session on device.

def Runtime.Autograd.Torch.Internal.SessionTape.size {α : Type} [TorchLean.Storage α] : SessionTape α → ℕ

Number of recorded nodes.

structure Runtime.Autograd.Torch.Internal.EagerSession (α : Type) [TorchLean.Storage α] : Type

The tape and side tables of eager execution.

paramsByLeaf links parameter leaves to their persistent storage. A reset drops the recording and advances refGeneration, but keeps the owner id, random counter, and backend selections.

• options : Config

  Session options controlling backend/device/kernel behavior.

• gradEnabled : Bool

  Whether newly recorded leaves track gradients. Inference sessions set this to false: forward values are still materialized so they can be read back, but parameter and input leaves are recorded with requiresGrad = false, and backend planning asks for no VJP.

• tape : IO.Ref (SessionTape α)

  The session's tape, on the device options selects.

• paramsByLeaf : IO.Ref (Std.HashMap ℕ (SomeParam α))

  Map from tape leaf ids to parameter objects, including frozen snapshot owners.

• parameterStorageByLeaf : IO.Ref (Std.HashMap ℕ (ParameterStorage α))

  Storage identities for this recording's parameter leaves; no tensor snapshots are cached.

• rngCounter : IO.Ref ℕ

  Number of seeded random tensors generated by this session.

• selectedBackends : IO.Ref (Array NN.Backend.AcceptedKernel)

  Accepted capsules already reported for this session, in first-use order.

• refOwner : ℕ

  Process-unique owner id for session references.

• refGeneration : IO.Ref ℕ

  Current recording generation for session references.

@[inline]

def Runtime.Autograd.Torch.Internal.EagerSession.keepTapeOnError {τ : Type} (t : τ) (result : Except String (τ × ℕ)) : τ × Except String ℕ

Pair the result of a tape constructor with the tape to store afterwards.

On success this is the extended tape. Tape constructors fail before they append, so on failure it is the tape the constructor started from. Call it inside the lambda passed to recordCpu or recordCuda, as fun t0 => keepTapeOnError t0 (Tape.op (t := t0) ...). There the constructor is inlined next to the match, the compiler sees that the failure branches are the only other uses of t0, and the final Array.push in addNode updates the node array in place.

def Runtime.Autograd.Torch.Internal.EagerSession.cpuTape {α : Type} [TorchLean.Storage α] (s : EagerSession α) : IO (Tape (Spec.SomeTensor α))

The session's tape, when the session runs on the CPU.

def Runtime.Autograd.Torch.Internal.EagerSession.cudaTape {α : Type} [TorchLean.Storage α] (s : EagerSession α) : IO (Tape LibTorch.SomeBuffer)

The session's tape, when the session runs on CUDA.

def Runtime.Autograd.Torch.Internal.EagerSession.setCudaTape {α : Type} [TorchLean.Storage α] (s : EagerSession α) (tape : Tape LibTorch.SomeBuffer) : IO Unit

Replace the tape of a CUDA session.

@[inline]

def Runtime.Autograd.Torch.Internal.EagerSession.recordCpu {α : Type} [TorchLean.Storage α] (s : EagerSession α) (op : Tape (Spec.SomeTensor α) → Tape (Spec.SomeTensor α) × Except String ℕ) : IO ℕ

Run a CPU tape step on the session tape and return the new node id.

The tape leaves the reference while op runs, and op returns the tape to store back together with the result (see keepTapeOnError). The tape is not used after op, so it stays unshared even when this function is not inlined, which happens when the calling IO action is passed around as a value, as in the backend dispatch.

@[inline]

def Runtime.Autograd.Torch.Internal.EagerSession.recordCpuPure {α : Type} [TorchLean.Storage α] (s : EagerSession α) (op : Tape (Spec.SomeTensor α) → Tape (Spec.SomeTensor α) × ℕ) : IO ℕ

recordCpu for constructors that cannot fail, such as leaves.

@[inline]

def Runtime.Autograd.Torch.Internal.EagerSession.recordCuda {α : Type} [TorchLean.Storage α] (s : EagerSession α) (op : Tape LibTorch.SomeBuffer → Tape LibTorch.SomeBuffer × Except String ℕ) : IO ℕ

recordCpu for the tape of a CUDA session.

@[inline]

def Runtime.Autograd.Torch.Internal.EagerSession.recordCudaPure {α : Type} [TorchLean.Storage α] (s : EagerSession α) (op : Tape LibTorch.SomeBuffer → Tape LibTorch.SomeBuffer × ℕ) : IO ℕ

recordCpuPure for the tape of a CUDA session.