Session #

TorchLean unified imperative session.

Session state #

A Session α is TorchLean's runtime analogue of a PyTorch "training loop environment". It:

owns a collection of leaf tensors (parameters and inputs),
records an eager computation on a tape or builds typed SSA graph data,
can run reverse-mode AD to produce gradients for all leaves, and
can apply simple optimizer steps (e.g. SGD) in a session-style workflow.

TorchLean exposes a single API with two execution modes selected at construction time:

.eager: a tape-backed runtime session (imperative autograd tape; useful for debugging and interactive examples),
.typedGraph: a session that records shape-indexed graph data while building the runtime tape, then executes via Runtime.Autograd.Torch.Internal.TypedGraphSession.

Both modes use the same Session API; each operation dispatches through Session.state.

Typical Training Loop (PyTorch Analogy) #

Think of the following mapping (approximately):

Session.param ~ create a torch.nn.Parameter (and later include it in a state_dict-like bundle).
Session.use ~ read a parameter as a tensor in the current recording phase.
Session.input ~ add a leaf tensor input (like feeding a batch tensor into the forward pass).
Session.resetTape ~ start a fresh recording phase (closest in spirit to optimizer.zero_grad() + new forward).
Session.backwardScalarDenseAll ~ loss.backward() (but returns gradients explicitly as an array).
Session.sgdStepAll ~ optimizer.step() (dense helper; higher-level training lives in NN.API.*).
Session.detach ~ tensor.detach() (cut the gradient edge at a value).

TorchLean does not store mutable .grad fields on each tensor ref; instead, gradients are returned explicitly (see grad, vjp, and the backward*DenseAll functions).

Non-Differentiable State (`NatRef`) #

NatRef stores the seed and counter used by the explicit random stream. Non-differentiable tensor inputs use the element-polymorphic data-input channel rather than a second session tensor type.

Deterministic RNG (Session-Level) #

RngState provides explicit, deterministic RNG state (closer to JAX PRNG keys than a global RNG). freshSeedIO is a convenience for sampling an initial seed at the IO boundary, while the core semantics remains seed-threaded and replayable.

Connection To TorchLean IR / Graph Execution #

In .typedGraph execution, the session records executable GraphData while building the tape. Each call to resetTape starts a new recording phase. Callers that need one reusable artifact should use Runtime.Autograd.Torch.TypedGraph directly; the high-level scalar trainer uses that artifact and records its loss graph once.

The type-level context checks graph shapes. Correctness of a stored JVP or VJP is a separate claim, proved only for operations connected to the proof-carrying Proofs.Autograd.Algebra.Node layer.

Practical note: the current .typedGraph implementation expects all leaves (tensor inputs/parameters and NatRefs) to be created before any op nodes are recorded. For portability, allocate leaves and initialize/split RNG up-front, then build the typed graph.

PyTorch References #

torch.autograd: https://pytorch.org/docs/stable/autograd.html
Tensor hooks (conceptual analogue of backwardDenseAllWithHook): https://pytorch.org/docs/stable/generated/torch.Tensor.register_hook.html

AD References #

This code follows the classic "tape / Wengert list" view of reverse-mode AD:

Andreas Griewank and Andrea Walther, Evaluating Derivatives, 2nd ed., 2008.
Seppo Linnainmaa, 1970 (reverse accumulation; precursor to modern backprop/autograd).