Session Shape and Index Operations

This file contains the session-level operations that preserve or rearrange tensor shape: activation helpers, reshapes, indexing, gathers, broadcasts, and reductions. Each operation dispatches through the same eager/typed graph session boundary as the lower-level tensor ops.

def Runtime.Autograd.Model.Session.relu {α : Type} [TorchLean.Storage α] (s : Session α) [Mul α] [Add α] [Zero α] [Max α] [BEq α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Rectified Linear Unit (ReLU) activation.

This is the pointwise nonlinearity $\operatorname{ReLU}(x)=\max(x,0)$, recorded as part of the session’s autograd graph.

PyTorch analogy: torch.relu(x) / torch.nn.functional.relu(x).

def Runtime.Autograd.Model.Session.sigmoid {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Sigmoid (logistic) activation, applied pointwise.

PyTorch analogy: torch.sigmoid(x).

def Runtime.Autograd.Model.Session.tanh {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Hyperbolic tangent activation, applied pointwise.

PyTorch analogy: torch.tanh(x).

def Runtime.Autograd.Model.Session.softmax {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis sh] (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Softmax along an explicitly selected tensor dimension.

def Runtime.Autograd.Model.Session.logSoftmax {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis sh] (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Stable log-softmax along an explicitly selected tensor dimension.

def Runtime.Autograd.Model.Session.softplus {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Softplus activation, applied pointwise: $\operatorname{softplus}(x)=\log(1+\exp x)$.

PyTorch analogy: torch.nn.functional.softplus(x).

def Runtime.Autograd.Model.Session.exp {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise exponential.

PyTorch analogy: torch.exp(x).

def Runtime.Autograd.Model.Session.sin {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise sine of angles in radians, using the session's eager or typed graph execution.

def Runtime.Autograd.Model.Session.cos {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise cosine, with gradients multiplied by -sin(x) on either execution path.

def Runtime.Autograd.Model.Session.log {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α sh)

Elementwise natural logarithm.

PyTorch analogy: torch.log(x).

If you need a total (always-defined) "log-like" surrogate without positivity side conditions, see safeLog.

def Runtime.Autograd.Model.Session.safeLog {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [Torch.TensorTransfer α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) (ε : α := Context.defaultEpsilon) : IO (Torch.TensorRef α sh)

Elementwise safe-log surrogate: $\operatorname{safeLog}(x;\varepsilon)=\log(\operatorname{softplus}(x)+\varepsilon)$.

We use this when we want something log-like but would rather not carry side conditions about inputs being strictly positive.

PyTorch analogy: torch.log(torch.nn.functional.softplus(x) + eps).

def Runtime.Autograd.Model.Session.sum {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α Spec.Shape.scalar)

Sum-reduce all elements of a tensor to a scalar.

PyTorch analogy: x.sum() (with no dim argument).

def Runtime.Autograd.Model.Session.flatten {α : Type} [TorchLean.Storage α] (s : Session α) [Inhabited α] [Zero α] {sh : Spec.Shape} (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α [sh.size])

Flatten a tensor to a 1D vector of length Spec.Shape.size sh.

PyTorch analogy: torch.flatten(x) or x.reshape(-1).

def Runtime.Autograd.Model.Session.reshape {α : Type} [TorchLean.Storage α] (s : Session α) [Inhabited α] [Zero α] {sh1 sh2 : Spec.Shape} (x : Torch.TensorRef α sh1) (h : sh1.size = sh2.size) : IO (Torch.TensorRef α sh2)

Reshape a tensor without changing the number of elements.

The proof h : Spec.Shape.size sh1 = Spec.Shape.size sh2 plays the role of PyTorch’s runtime check performed by reshape/view.

PyTorch analogy: x.reshape(new_shape) (when the element count matches).

def Runtime.Autograd.Model.Session.swapAdjacentAtDepth {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (depth : ℕ) (x : Torch.TensorRef α sh) : IO (Torch.TensorRef α (sh.swapAdjacentAtDepth depth))

Generic "swap adjacent axes" view operation.

This is a shape-driven permutation helper used in some attention/transformer code.

def Runtime.Autograd.Model.Session.broadcastTo {α : Type} [TorchLean.Storage α] (s : Session α) [Inhabited α] [Add α] [Zero α] {sh1 sh2 : Spec.Shape} (cb : sh1.CanBroadcastTo sh2) (x : Torch.TensorRef α sh1) : IO (Torch.TensorRef α sh2)

Broadcast a tensor to a larger shape (dispatches by execution mode).

def Runtime.Autograd.Model.Session.reduceSum {α : Type} [TorchLean.Storage α] (s : Session α) [Add α] [Zero α] [Inhabited α] {sh : Spec.Shape} (axis : ℕ) (x : Torch.TensorRef α sh) [valid : Spec.Shape.HasNonemptyAxis axis sh] [wf : sh.WellFormed] : IO (Torch.TensorRef α (TorchLean.Tensor.shapeAfterSum sh axis))

Reduce-sum along an axis (dispatches by execution mode).

def Runtime.Autograd.Model.Session.reduceMean {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {sh : Spec.Shape} (axis : ℕ) (x : Torch.TensorRef α sh) [valid : Spec.Shape.HasNonemptyAxis axis sh] [wf : sh.WellFormed] : IO (Torch.TensorRef α (TorchLean.Tensor.shapeAfterSum sh axis))

Reduce-mean along an axis (dispatches by execution mode).

def Runtime.Autograd.Model.Session.select {α : Type} [TorchLean.Storage α] (s : Session α) [Zero α] {shape : Spec.Shape} (axis : ℕ) (x : Torch.TensorRef α shape) [Spec.Shape.AxisInBounds axis shape] (index : Fin (shape.axisSize axis)) : IO (Torch.TensorRef α (shape.eraseAxis axis))

Select one bounded coordinate from an arbitrary tensor axis.

def Runtime.Autograd.Model.Session.indexSelect {α : Type} [TorchLean.Storage α] (s : Session α) [Add α] [Zero α] {shape : Spec.Shape} (axis count : ℕ) (x : Torch.TensorRef α shape) [Spec.Shape.AxisInBounds axis shape] (indices : TorchLean.Tensor (Fin (shape.axisSize axis)) [count]) : IO (Torch.TensorRef α (shape.replaceAxis axis count))

Select several bounded coordinates from an arbitrary tensor axis.

def Runtime.Autograd.Model.Session.scatterAdd {α : Type} [TorchLean.Storage α] (s : Session α) [Add α] [Zero α] {shape : Spec.Shape} (axis count : ℕ) (base : Torch.TensorRef α shape) [Spec.Shape.AxisInBounds axis shape] (source : Torch.TensorRef α (shape.replaceAxis axis count)) (indices : TorchLean.Tensor (Fin (shape.axisSize axis)) [count]) : IO (Torch.TensorRef α shape)

Add source slices into an arbitrary tensor axis at bounded coordinates.