Eager Tensor Operations

PyTorch-style tensor operations backed by the eager CPU/CUDA tapes. These wrappers record runtime nodes, dispatch CUDA kernels when requested, and preserve the typed TensorRef surface.

Shape and reduction operations

def Runtime.Autograd.Torch.Internal.EagerSession.sum {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α Spec.Shape.scalar)

Sum-reduce all elements to a scalar. PyTorch: x.sum().

def Runtime.Autograd.Torch.Internal.EagerSession.flatten {α : Type} (s : EagerSession α) [Inhabited α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α (Spec.Shape.dim sh.size Spec.Shape.scalar))

Flatten a tensor to a 1D vector. PyTorch: torch.flatten.

def Runtime.Autograd.Torch.Internal.EagerSession.reshape {α : Type} (s : EagerSession α) [Inhabited α] [DecidableEq Spec.Shape] {sh1 sh2 : Spec.Shape} (x : TensorRef α sh1) (h : sh1.size = sh2.size) : IO (TensorRef α sh2)

Reshape a tensor while preserving total number of elements.

PyTorch comparison: torch.reshape / view (when valid).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose2d {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {m n : ℕ} (x : TensorRef α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : IO (TensorRef α (Spec.Shape.dim n (Spec.Shape.dim m Spec.Shape.scalar)))

Transpose a 2D matrix. PyTorch: x.t() / x.transpose(0,1).

def Runtime.Autograd.Torch.Internal.EagerSession.swapAdjacentAtDepth {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {sh : Spec.Shape} (depth : ℕ) (x : TensorRef α sh) : IO (TensorRef α (sh.swapAdjacentAtDepth depth))

Swap two adjacent axes at a given depth. PyTorch analogue: x.transpose(dim, dim+1).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose3dFirstToLast {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {a b c : ℕ} (x : TensorRef α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim b (Spec.Shape.dim c (Spec.Shape.dim a Spec.Shape.scalar))))

Permute a 3D tensor (a,b,c) → (b,c,a). PyTorch: x.permute(1,2,0).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose3dLastToFirst {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {a b c : ℕ} (x : TensorRef α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim c (Spec.Shape.dim a (Spec.Shape.dim b Spec.Shape.scalar))))

Permute a 3D tensor (a,b,c) → (c,a,b). PyTorch: x.permute(2,0,1).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose3dLastTwo {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {a b c : ℕ} (x : TensorRef α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim a (Spec.Shape.dim c (Spec.Shape.dim b Spec.Shape.scalar))))

Swap the last two axes of a 3D tensor (a,b,c) → (a,c,b). PyTorch: x.transpose(1,2).

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

Broadcast a tensor to a larger shape. PyTorch: implicit broadcasting / expand.

def Runtime.Autograd.Torch.Internal.EagerSession.reduceSum {α : Type} (s : EagerSession α) [Add α] [Zero α] [Inhabited α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis sh] [wf : sh.WellFormed] (x : TensorRef α sh) : IO (TensorRef α (Spec.Tensor.shapeAfterSum sh axis))

Sum-reduce along axis. PyTorch: torch.sum(x, dim=axis).

def Runtime.Autograd.Torch.Internal.EagerSession.reduceMean {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis sh] [wf : sh.WellFormed] (x : TensorRef α sh) : IO (TensorRef α (Spec.Tensor.shapeAfterSum sh axis))

Mean-reduce along axis. PyTorch: torch.mean(x, dim=axis).