Session Tensor Operations

This file is the shared execution-mode dispatch layer for elementary tensor operations. It keeps the public Session API uniform while routing each operation to either the eager tape or the typed SSA graph session.

def Runtime.Autograd.Model.Session.add {α : Type} [TorchLean.Storage α] (s : Session α) [Add α] [Zero α] {sh : Spec.Shape} (a b : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise addition (dispatches by execution mode).

def Runtime.Autograd.Model.Session.sub {α : Type} [TorchLean.Storage α] (s : Session α) [Sub α] [Add α] [Zero α] {sh : Spec.Shape} (a b : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise subtraction (dispatches by execution mode).

def Runtime.Autograd.Model.Session.mul {α : Type} [TorchLean.Storage α] (s : Session α) [Mul α] [Add α] [Zero α] {sh : Spec.Shape} (a b : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise multiplication (dispatches by execution mode).

def Runtime.Autograd.Model.Session.scale {α : Type} [TorchLean.Storage α] (s : Session α) [Mul α] [Add α] [Zero α] [Torch.TensorTransfer α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) (c : α) : IO (Torch.TensorRef α sh)

Scale a tensor by a scalar constant c (elementwise).

PyTorch analogy: x * c or torch.mul(x, c).

def Runtime.Autograd.Model.Session.dropout {α : Type} [TorchLean.Storage α] [Context α] [Torch.TensorTransfer α] (s : Session α) (rng : RngState) {sh : Spec.Shape} (x : Torch.TensorRef α sh) (p : α) (train : Bool := true) : IO (Torch.TensorRef α sh)

Dropout implemented as a Session-level derived op.

In training mode, this records $$ y=\frac{x\odot\mathtt{mask}}{1-p}, $$

where $\mathtt{mask}\in\{0,1\}$ is generated by the Torch primitive bernoulliMask.

RNG Semantics (JAX-style / functional RNG)

Randomness is a deterministic function of explicit session state:

• we read (seed, counter) from the session-level RngState,
• derive a per-call opSeed deterministically, and
• advance the RngState (update seed, increment counter).

Important: the bernoulli_mask op also mixes in a backend-internal counter (approximately: current tape size / node index). So even with the same RngState, changing the surrounding graph structure can change the exact samples. This is still fully deterministic for a fixed graph.

In evaluation mode (train=false) and at p = 0, this is the identity. At p = 1, training returns zero. Other values must satisfy $0 < p < 1$.

def Runtime.Autograd.Model.Session.abs {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise absolute value (dispatches by execution mode).

def Runtime.Autograd.Model.Session.sqrt {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise square root (dispatches by execution mode).

def Runtime.Autograd.Model.Session.clamp {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [Torch.TensorTransfer α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) (minVal maxVal : α) : IO (Torch.TensorRef α sh)

Elementwise clamp to [minVal, maxVal] (dispatches by execution mode).

def Runtime.Autograd.Model.Session.max {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {sh : Spec.Shape} (a b : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise maximum (dispatches by execution mode).

def Runtime.Autograd.Model.Session.min {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {sh : Spec.Shape} (a b : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise minimum (dispatches by execution mode).

def Runtime.Autograd.Model.Session.matmul {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {batchA batchB batch : Spec.Shape} {m n p : ℕ} [broadcastA : batchA.BroadcastTo batch] [broadcastB : batchB.BroadcastTo batch] (a : Torch.TensorRef α (batchA.concat [m, n])) (b : Torch.TensorRef α (batchB.concat [n, p])) : IO (Torch.TensorRef α (batch.concat [m, p]))

Matrix multiplication with broadcasted batch prefixes.

def Runtime.Autograd.Model.Session.concatLeadingAxis {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {n m : ℕ} {sh : Spec.Shape} (a : Torch.TensorRef α (Spec.Shape.dim n sh)) (b : Torch.TensorRef α (Spec.Shape.dim m sh)) : IO (Torch.TensorRef α (Spec.Shape.dim (n + m) sh))

Concatenate along the outermost dimension (dimension 0) (dispatches to eager vs typed graph backend).

def Runtime.Autograd.Model.Session.sliceLeadingAxisRange {α : Type} [TorchLean.Storage α] (s : Session α) [Zero α] {n : ℕ} {sh : Spec.Shape} (x : Torch.TensorRef α (Spec.Shape.dim n sh)) (start len : ℕ) (h : start + len ≤ n) : IO (Torch.TensorRef α (Spec.Shape.dim len sh))

Slice a contiguous [start, start+len) range from dimension 0.

PyTorch analogy: x[start:start+len] for the first dimension.

def Runtime.Autograd.Model.Session.maxPool {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {d channels : ℕ} {spatial kernel stride padding : TorchLean.Tensor ℕ [d]} (x : Torch.TensorRef α (Spec.Shape.ofList (channels :: spatial.to (List ℕ)))) : IO (Torch.TensorRef α (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial kernel stride padding).to (List ℕ))))

Apply max pooling over an arbitrary number of spatial axes.

def Runtime.Autograd.Model.Session.smoothMaxPool {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [DecidableEq α] [Torch.TensorTransfer α] {d channels : ℕ} {spatial kernel stride padding : TorchLean.Tensor ℕ [d]} (x : Torch.TensorRef α (Spec.Shape.ofList (channels :: spatial.to (List ℕ)))) (beta : α) : IO (Torch.TensorRef α (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial kernel stride padding).to (List ℕ))))

Apply smooth max pooling over an arbitrary number of spatial axes.

def Runtime.Autograd.Model.Session.avgPool {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {d channels : ℕ} {spatial kernel stride padding : TorchLean.Tensor ℕ [d]} (x : Torch.TensorRef α (Spec.Shape.ofList (channels :: spatial.to (List ℕ)))) : IO (Torch.TensorRef α (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial kernel stride padding).to (List ℕ))))

Apply average pooling over an arbitrary number of spatial axes.