Runtime context

This module defines the runtime environment for executing TorchLean computations.

The main challenge is that tensors are dependently typed by their Shape, but at runtime we want to store a heterogeneous map from names to values. We solve this by storing an existential wrapper AnyTensor that carries the Shape alongside the tensor. Lookup functions then check the shape and use the resulting equality proof to cast the stored tensor to the requested shape.

Reading map

• AnyTensor is the shape-erased wrapper used by the registries.
• RuntimeContext is the mutable-looking record that stores values, gradients, and a fresh-id counter.
• register_variable / get_variable are the core value registry operations.
• register_gradient / get_gradient do the same on the gradient side.

structure Runtime.AnyTensor (α : Type) : Type

Existential wrapper for tensors of arbitrary shape.

This is the core trick that lets a runtime registry store tensors of different shapes in one container.

• s : Spec.Shape

  The shape carried alongside the tensor value.

• t : Spec.Tensor α self.s

  The tensor value, indexed by s.

structure Runtime.RuntimeContext (α : Type) : Type

Runtime context for tracking named variables and their gradients.

We keep two registries:

• var_registry for values, and
• gradients for accumulated gradients (typically produced by backprop).

• var_registry : List (String × AnyTensor α)

  Registry of named variable values.

  Implementation note: this is a simple list, so lookup uses the first match (most-recent wins).

• gradients : List (String × AnyTensor α)

  Registry of named gradients.

  As with var_registry, this is a simple list; callers are responsible for deciding whether to prepend new gradients or update/accumulate existing ones.

• next_id : ℕ

  Fresh id counter for allocating new variables/parameters.

  This is incremented by register_variable and used by higher-level runtime layers.

def Runtime.emptyContext {α : Type} : RuntimeContext α

An empty context with no variables and no gradients.

def Runtime.registerVariable {α : Type} {s : Spec.Shape} (ctx : RuntimeContext α) (name : String) (value : Spec.Tensor α s) : RuntimeContext α

Register a new variable in the context.

The value is wrapped as an AnyTensor so we can store it in a heterogeneous registry.

def Runtime.getVariable {α : Type} {s : Spec.Shape} [DecidableEq Spec.Shape] (ctx : RuntimeContext α) (name : String) : Option (Spec.Tensor α s)

Lookup a variable by name and requested shape.

If the name exists and the stored Shape matches s, we cast the stored tensor from Tensor α any_tensor.s to Tensor α s using the equality proof any_tensor.s = s.

def Runtime.setVariable {α : Type} {s : Spec.Shape} (ctx : RuntimeContext α) (name : String) (value : Spec.Tensor α s) : RuntimeContext α

Update the value of an existing variable (or do nothing if the name is absent).

This performs a simple list-map update of the registry.

def Runtime.registerGradient {α : Type} {s : Spec.Shape} (ctx : RuntimeContext α) (name : String) (grad : Spec.Tensor α s) : RuntimeContext α

Register (prepend) a gradient entry in the context.

If you want to accumulate gradients under the same name, do that at a higher layer before calling this helper.

def Runtime.getGradient {α : Type} {s : Spec.Shape} [DecidableEq Spec.Shape] (ctx : RuntimeContext α) (name : String) : Option (Spec.Tensor α s)

Lookup a gradient by name and requested shape (with the same casting trick as get_variable).

def Runtime.setGradient {α : Type} {s : Spec.Shape} (ctx : RuntimeContext α) (name : String) (grad : Spec.Tensor α s) : RuntimeContext α

Update a gradient entry in the context (simple list-map update).

def Runtime.clearGradients {α : Type} (ctx : RuntimeContext α) : RuntimeContext α

Remove all gradient entries (analogue of PyTorch optimizer.zero_grad()).

def Runtime.hasVariable {α : Type} (ctx : RuntimeContext α) (name : String) : Bool

Return true iff the context contains a variable named name.

def Runtime.variableNames {α : Type} (ctx : RuntimeContext α) : List String

List all variable names stored in ctx.

def Runtime.gradientNames {α : Type} (ctx : RuntimeContext α) : List String

List all gradient names stored in ctx.

def Runtime.contextSize {α : Type} (ctx : RuntimeContext α) : ℕ

Number of registered variables.

def Runtime.gradientCount {α : Type} (ctx : RuntimeContext α) : ℕ

Number of registered gradient entries.

def Runtime.isValidContext {α : Type} (ctx : RuntimeContext α) : Bool

Check a simple invariant: every gradient entry refers to an existing variable name.

This does not check shapes; it only checks name presence.

def Runtime.contextToString {α : Type} [ToString α] (ctx : RuntimeContext α) : String

Render the context contents as a string (for debugging).