Torch Core

PyTorch-style imperative front-end (eager runtime).

This wraps the eager runtime tape (Runtime.Autograd.Tape) behind an IO.Ref so user code can look closer to PyTorch:

• create parameters and inputs as objects;
• call ops as methods (no explicit tape threading);
• run backward and optionally apply SGD updates.

This is purely a convenience layer; correctness/proof connections live elsewhere (e.g. Runtime.Autograd.Compiled / Proofs.Autograd.Algebra.Graph).

References:

• PyTorch torch.autograd docs: https://pytorch.org/docs/stable/autograd.html
• PyTorch "Autograd mechanics": https://pytorch.org/docs/stable/notes/autograd.html

@[reducible, inline]

abbrev Runtime.Autograd.Torch.TList (α : Type) (ss : List Spec.Shape) : Type

TList is the dependently-typed heterogeneous tensor list used by the proved IR.

We re-export it here under the Torch front-end namespace because many user-facing helpers (trainer APIs, parameter packs, etc.) are naturally expressed as TLists.

inductive Runtime.Autograd.Torch.Backend : Type

Execution backend for the Torch-style front-end.

• .eager: build and execute a runtime tape directly (imperative, PyTorch-like).
• .compiled: record typed IR and run a compiled tape (proof-friendly path, see Torch.LinkedSession / TorchLean.Session).

This is intentionally not a CUDA Graph selector. CUDA is controlled by Options.device on the eager backend; CUDA Graph capture/replay will require a distinct persistent-buffer backend.

• eager : Backend
• compiled : Backend

@[implicit_reducible]

instance Runtime.Autograd.Torch.instReprBackend : Repr Backend

def Runtime.Autograd.Torch.instReprBackend.repr : Backend → ℕ → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.Torch.instDecidableEqBackend : DecidableEq Backend

inductive Runtime.Autograd.Torch.Device : Type

Execution device selector (PyTorch comparison: cpu vs cuda).

Current scope:

• eager backend only: selects whether the hidden tape is CPU (Runtime.Autograd.Tape) or CUDA (Runtime.Autograd.Cuda.Tape),
• compiled backend is unchanged (proof semantics / typed IR are device-agnostic).

• cpu : Device
• cuda : Device

def Runtime.Autograd.Torch.instReprDevice.repr : Device → ℕ → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.Torch.instReprDevice : Repr Device

@[implicit_reducible]

instance Runtime.Autograd.Torch.instDecidableEqDevice : DecidableEq Device

structure Runtime.Autograd.Torch.Options : Type

Options controlling the behavior of the Torch-style front-end.

PyTorch comparison: these are roughly "session/global" settings (e.g. default requires_grad, and runtime-only performance toggles).

• backend : Backend

  Execution backend selection.

• requiresGradByDefault : Bool

  Default requires_grad value for newly created parameters/inputs when a caller omits it.

• seed : ℕ

  Global deterministic seed for demo-style randomness.

  TorchLean keeps the semantic core pure and seed-threaded (JAX-style), so this is best understood as a convenient default seed knob that user code can thread into:

  • model initialization (per-layer init keys),
  • dataset shuffles / sampling,
  • and session-level RNG state (dropout, etc.).

  PyTorch analogue: torch.manual_seed(seed).

• fastKernels : Bool

  Enable runtime-only fast kernels for a few hot ops in the eager backend.

  This is an execution/performance flag; it is not used by the proof-linked compilation path.

• fastGpuMatmulPrecision : FastKernels.GpuMatmulPrecision

  GPU precision for fast-kernel matmul over Lean Float tensors.

  .fp32 matches the eager CUDA buffer stack. .fp64 selects the double-precision DGEMM path for matmul-only Float workloads that intentionally want double precision on GPU.

• useGpu : Bool

  Eager execution on CUDA.

  When true and backend = .eager, the eager session uses the CUDA tape (Runtime.Autograd.Cuda.Tape) and Torch ops must route to CUDA implementations (no implicit CPU fallback).

  Compiled backend behavior is unchanged.

• strictCuda : Bool

  Strict CUDA mode (eager backend only).

  Note: in CUDA eager mode TorchLean does not fall back to CPU per-op; missing CUDA ops always throw. This flag is retained for API compatibility.

def Runtime.Autograd.Torch.instReprOptions.repr : Options → ℕ → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.Torch.instReprOptions : Repr Options

def Runtime.Autograd.Torch.Options.device (opts : Options) : Device

Read the device selector corresponding to useGpu.

def Runtime.Autograd.Torch.Options.toDevice (opts : Options) (d : Device) : Options

Set the device selector.

structure Runtime.Autograd.Torch.TensorRef (α : Type) (s : Spec.Shape) : Type

Opaque handle to a tensor value in the current session/tape.

This is the TorchLean analogue of a PyTorch Tensor object whose "identity" is a node/leaf id in the autograd tape. The phantom shape index s makes shape mismatches explicit at compile time.

• id : ℕ

  Node/leaf identifier in the owning session tape.

def Runtime.Autograd.Torch.instReprTensorRef.repr {α✝ : Type} {s✝ : Spec.Shape} [Repr α✝] : TensorRef α✝ s✝ → ℕ → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.Torch.instReprTensorRef {α✝ : Type} {s✝ : Spec.Shape} [Repr α✝] : Repr (TensorRef α✝ s✝)

structure Runtime.Autograd.Torch.NatRef : Type

Handle to a Nat stored in the session's non-differentiable environment.

This is used to model index-like inputs (class labels, gather indices, etc.) which should not receive gradients.

• id : ℕ

  Index into the session's non-differentiable Nat environment.

@[implicit_reducible]

instance Runtime.Autograd.Torch.instReprNatRef : Repr NatRef

def Runtime.Autograd.Torch.instReprNatRef.repr : NatRef → ℕ → Std.Format

structure Runtime.Autograd.Torch.NatVecRef (k : ℕ) : Type

Handle to a contiguous block of k Nats in the session's non-differentiable environment.

• start : ℕ

  Start offset of the contiguous k-element block in the session's Nat environment.

def Runtime.Autograd.Torch.instReprNatVecRef.repr {k✝ : ℕ} : NatVecRef k✝ → ℕ → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.Torch.instReprNatVecRef {k✝ : ℕ} : Repr (NatVecRef k✝)

structure Runtime.Autograd.Torch.Param (α : Type) (s : Spec.Shape) : Type

Trainable parameter: a mutable tensor value plus metadata.

PyTorch comparison: analogous to torch.nn.Parameter, except the parameter becomes part of the autograd graph only when you use it in a particular session/tape.

• name : Option String

  Optional user-facing name for logging/debugging.

• value : IO.Ref (Spec.Tensor α s)

  Value at the current point.

• cudaValue : IO.Ref (Option Cuda.AnyBuffer)

  Optional CUDA-resident mirror of value.

  The eager CUDA trainer uses this as a lightweight persistent-parameter cache: repeated forward passes can reuse the device buffer instead of uploading the host tensor every step. The host value remains the public source for CPU/proof-oriented APIs and is synchronized on explicit readback.

• hostCurrent : IO.Ref Bool

  Whether value is known to match cudaValue.

  CUDA optimizer steps mark this false after updating only the device mirror. Public parameter readback synchronizes and flips it back to true.

• requiresGrad : Bool

  Whether this parameter receives accumulated gradients and optimizer updates.

structure Runtime.Autograd.Torch.AnyParam (α : Type) : Type

Type-erased parameter wrapper.

This exists so session code can store heterogeneous parameter shapes in a single HashMap keyed by leaf id (used for SGD updates).

• s : Spec.Shape

  Runtime shape of the erased parameter.

• requiresGrad : Bool

  Whether the underlying parameter receives optimizer updates.

• get : IO (AnyTensor α)

  Read the current parameter value with its runtime shape.

• set : AnyTensor α → IO Unit

  Overwrite the current parameter value, checking shape at the call site.

• setCuda : Cuda.AnyBuffer → IO Unit

  Store a CUDA buffer mirror without forcing an immediate host download.

def Runtime.Autograd.Torch.AnyParam.releaseCachedCudaValue {α : Type} {s : Spec.Shape} (p : Param α s) : IO Unit

Release a cached CUDA mirror, if one exists.

CUDA buffers are external objects whose native finalizer tolerates repeated cleanup attempts, but long training loops should not wait for GC pressure before returning replaced parameter mirrors to the allocator.

def Runtime.Autograd.Torch.AnyParam.ofParam {α : Type} {s : Spec.Shape} (p : Param α s) : AnyParam α

Package a typed Param α s as an AnyParam α, checking shape on set.

This is the bridge that allows generic optimizers/update routines to operate over heterogeneous parameter packs.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.okOrThrow {α : Type} : Result α → IO α

Convenience: throw IO.userError on a .error result.

Eager backend internals

The eager backend for the backend-generic Ops interface needs a small tape-backed session to thread an IO.Ref to the runtime tape.

This is intentionally kept under Torch.Internal.*; the public session-style API is Runtime.Autograd.TorchLean.Session.

CUDA Bridge (Upload/Download)

The CUDA eager tape stores float32 device buffers (Runtime.Autograd.Cuda.Buffer) paired with a runtime Shape (Runtime.Autograd.Cuda.AnyBuffer).

The Torch eager front-end still presents the spec-level Tensor α s API, so in CUDA mode we need:

• upload: Tensor α s -> Cuda.AnyBuffer (float32, contiguous, row-major)
• download: Cuda.AnyBuffer -> Runtime.AnyTensor α

The helper namespace gives CUDA bridge conversions stable call sites and a clear boundary.

class Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv (α : Type) : Type

Conversions required by the eager CUDA tape path.

• toAnyBuffer {s : Spec.Shape} : Spec.Tensor α s → IO Cuda.AnyBuffer

  Upload a spec tensor to a CUDA AnyBuffer (float32).

• ofAnyBuffer : Cuda.AnyBuffer → IO (AnyTensor α)

  Download a CUDA AnyBuffer to a runtime AnyTensor (shape-erased).

• toFloat : α → IO Float

  Convert a scalar constant to a host Float for CUDA kernels (e.g. scale, axpy).

Float implementation

@[implicit_reducible]

instance Runtime.Autograd.Torch.Internal.CudaBridge.instTensorConvFloat : TensorConv Float

Float CUDA conversions: upload/download via row-major FloatArray.

@[implicit_reducible, instance 10]

instance Runtime.Autograd.Torch.Internal.CudaBridge.instTensorConv (α : Type) : TensorConv α

Generic CPU-preserving fallback for scalar types without a CUDA wire-format bridge.

Many TorchLean sessions are scalar-polymorphic on CPU, while the eager CUDA tape stores float32 buffers. The fallback keeps CPU execution available for proof-oriented scalar backends and fails loudly if a CUDA-only conversion is actually requested. Add a higher-priority TensorConv α instance for scalar types that have a deliberate float32 wire representation.

Shape helpers for CUDA kernels

def Runtime.Autograd.Torch.Internal.CudaBridge.dimsArray (s : Spec.Shape) : Array ℕ

Runtime dimension list as an Array Nat (outermost-first).

def Runtime.Autograd.Torch.Internal.CudaBridge.axisMapArray {s₁ s₂ : Spec.Shape} (cb : s₁.CanBroadcastTo s₂) : Array ℕ

axisMap as an array.

def Runtime.Autograd.Torch.Internal.syncParamCudaToHost {α : Type} [CudaBridge.TensorConv α] {sh : Spec.Shape} [DecidableEq Spec.Shape] (p : Param α sh) : IO Unit

Synchronize a CUDA-updated parameter back to its host tensor, if needed.

This is deliberately explicit. Training hot paths keep parameters resident on device; public readback APIs call this helper before exposing parameter tensors to the Lean side.

def Runtime.Autograd.Torch.Internal.setParamCudaValue {α : Type} {sh : Spec.Shape} (p : Param α sh) (any : Cuda.AnyBuffer) : IO Unit

Store/update the CUDA mirror of a parameter and mark the host tensor stale.

def Runtime.Autograd.Torch.Internal.setParamHostValue {α : Type} {sh : Spec.Shape} (p : Param α sh) (v : Spec.Tensor α sh) : IO Unit

Overwrite a host parameter value and invalidate any stale CUDA mirror.

structure Runtime.Autograd.Torch.Internal.EagerSession (α : Type) : Type

Internal eager session: a mutable runtime tape plus side tables.

This is the state needed to offer a PyTorch-like API where "tensors" are opaque references and ops mutate a hidden tape stored in an IO.Ref.

Notes:

• tape stores values and backward closures (Runtime.Autograd.Tape).
• paramsByLeaf remembers which tape leaf ids correspond to trainable parameters (for SGD).
• nats stores non-differentiable Nat inputs used for indexing-like operations.

• opts : Options

  Session options controlling backend/device/kernel behavior.

• tape : IO.Ref (Tape α)

  CPU eager tape used when opts.useGpu = false.

• cudaTape : IO.Ref Cuda.Tape

  CUDA eager tape used when opts.useGpu = true.

• paramsByLeaf : IO.Ref (Std.HashMap ℕ (AnyParam α))

  Map from tape leaf ids to trainable parameter objects.

• nats : IO.Ref (Array ℕ)

  Non-differentiable integer inputs for dynamic indexing operations.

def Runtime.Autograd.Torch.Internal.EagerSession.new {α : Type} (opts : Options := { }) : IO (EagerSession α)

Allocate a fresh eager session with an empty tape and empty side tables.

def Runtime.Autograd.Torch.Internal.EagerSession.releaseCudaBuffer (b : Cuda.Buffer) : IO Unit

Force-free a CUDA buffer allocation; the external finalizer is safe to call twice.

def Runtime.Autograd.Torch.Internal.EagerSession.releaseCudaAnyBuffer (b : Cuda.AnyBuffer) : IO Unit

Force-release a shape-erased CUDA buffer.

def Runtime.Autograd.Torch.Internal.EagerSession.releaseCudaTapeNonParamValues {α : Type} (s : EagerSession α) : IO Unit

Release current CUDA tape values that are not persistent parameter mirrors.

Eager CUDA training creates many temporary buffers per step. Relying only on external-object finalizers can produce high transient memory pressure in long runs, so reset/step paths call this before discarding the current tape snapshot.

def Runtime.Autograd.Torch.Internal.EagerSession.releaseCudaTapeAfterOptimizerStep {α : Type} (s : EagerSession α) : IO Unit

Release CUDA tape values after an optimizer step.

Unlike releaseCudaTapeNonParamValues, this may release trainable parameter leaf buffers too. In a CUDA optimizer step, trainable parameters have already been written to fresh persistent mirrors, so the leaf buffers from the just-consumed tape are stale. Non-trainable parameter leaves still are their persistent mirrors, so we keep those cached.

def Runtime.Autograd.Torch.Internal.EagerSession.releaseCudaAnyBufferArray (xs : Array Cuda.AnyBuffer) : IO Unit

Release a dense CUDA gradient array after an optimizer has consumed it.

def Runtime.Autograd.Torch.Internal.EagerSession.collectCudaAllocator : IO Unit

Ask the native allocator to return/free pages after a large CUDA eager step.

def Runtime.Autograd.Torch.Internal.EagerSession.resetTape {α : Type} (s : EagerSession α) : IO Unit

Reset the tape and side tables.

PyTorch comparison: like starting a fresh forward pass where the autograd graph is discarded.

def Runtime.Autograd.Torch.Internal.EagerSession.param {α : Type} (s : EagerSession α) {sh : Spec.Shape} (init : Spec.Tensor α sh) (name : Option String := none) (requiresGrad : Option Bool := none) : IO (Param α sh)

Create a mutable parameter object (not yet on the tape).

To record this parameter on the session tape, call use, which reads the parameter and records it as a leaf.

def Runtime.Autograd.Torch.Internal.EagerSession.getValue {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (x : TensorRef α sh) : IO (Spec.Tensor α sh)

Read back the concrete tensor value stored at a TensorRef.

This is a dynamic check: we ensure the id exists on the tape and the stored shape matches sh.

def Runtime.Autograd.Torch.Internal.EagerSession.const {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (v : Spec.Tensor α sh) (name : Option String := none) : IO (TensorRef α sh)

Record a constant leaf (non-differentiable) on the tape.

PyTorch comparison: like constructing a tensor with requires_grad=False.

def Runtime.Autograd.Torch.Internal.EagerSession.detach {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (x : TensorRef α sh) (name : Option String := none) : IO (TensorRef α sh)

Stop-gradient boundary.

Forward semantics: identity (detach(x) = x). Backward semantics: no gradient flows to x.

def Runtime.Autograd.Torch.Internal.EagerSession.randUniform {α : Type} [Context α] [CudaBridge.TensorConv α] (s : EagerSession α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (seed : ℕ) (name : Option String := none) : IO (TensorRef α sh)

Deterministic U[0,1) tensor generator (seeded).

def Runtime.Autograd.Torch.Internal.EagerSession.bernoulliMask {α : Type} [Context α] [CudaBridge.TensorConv α] (s : EagerSession α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (keepProb : TensorRef α Spec.Shape.scalar) (seed : ℕ) (name : Option String := none) : IO (TensorRef α sh)

Deterministic {0,1} mask generator (seeded) with a scalar keep-probability input.

def Runtime.Autograd.Torch.Internal.EagerSession.use {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (p : Param α sh) : IO (TensorRef α sh)

Use a parameter in the tape by recording its current value as a leaf.

The returned TensorRef is the handle you pass to ops. The leaf id is stored in paramsByLeaf so optimizer steps (e.g. SGD) can update parameters after backward. PyTorch comparison: like using a torch.nn.Parameter in a forward pass (it becomes a leaf in the autograd graph).

def Runtime.Autograd.Torch.Internal.EagerSession.input {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (v : Spec.Tensor α sh) (name : Option String := none) (requiresGrad : Bool := false) : IO (TensorRef α sh)

Record an external input tensor as a leaf on the tape.

PyTorch comparison: like introducing a tensor into the autograd graph with a chosen requires_grad flag.

def Runtime.Autograd.Torch.Internal.EagerSession.inputNat {α : Type} (s : EagerSession α) (v : ℕ) : IO NatRef

Record a non-differentiable Nat input in the session environment.

This supports ops that depend on indices/labels that should not receive gradients.

def Runtime.Autograd.Torch.Internal.EagerSession.getNat {α : Type} (s : EagerSession α) (r : NatRef) : IO ℕ

Read a previously recorded NatRef.

def Runtime.Autograd.Torch.Internal.EagerSession.setNat {α : Type} (s : EagerSession α) (r : NatRef) (v : ℕ) : IO Unit

Overwrite a previously recorded NatRef.

def Runtime.Autograd.Torch.Internal.EagerSession.natVecToArray {k : ℕ} (v : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : Array ℕ

Convert a Tensor Nat (.dim k .scalar) to an Array Nat.

Used to stage NatVecRef values into the session environment.

def Runtime.Autograd.Torch.Internal.EagerSession.inputNatVec {α : Type} {k : ℕ} (s : EagerSession α) (v : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO (NatVecRef k)

Record a non-differentiable vector of Nats in the session environment.

Returns a NatVecRef k pointing to the stored block.

def Runtime.Autograd.Torch.Internal.EagerSession.getNatVec {α : Type} {k : ℕ} (s : EagerSession α) (r : NatVecRef k) : IO (Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar))

Read back the vector stored at NatVecRef k.

def Runtime.Autograd.Torch.Internal.EagerSession.setNatVec {α : Type} {k : ℕ} (s : EagerSession α) (r : NatVecRef k) (v : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO Unit

Overwrite the stored vector at NatVecRef k.

Tensor ops (eager tape wrappers)

The following definitions are thin wrappers around Runtime.Autograd.Tape.* primitives. Each one:

• reads the current tape from s.tape,
• appends a new node/leaf via a Tape.* constructor,
• writes the updated tape back, and
• returns a fresh TensorRef pointing to the new node id.

PyTorch comparison: this is the standard eager autograd mechanism (a dynamic tape of ops).

def Runtime.Autograd.Torch.Internal.EagerSession.dispatchCudaOpt {α β : Type} (s : EagerSession α) (opName : String) (cpu : IO β) (cuda : IO (Option β)) : IO β

Dispatch an eager op with optional CUDA support.

When Options.device = .cuda, any op whose CUDA implementation returns none will throw.

TorchLean's CUDA eager mode is intentionally "no per-op CPU fallback": either the op is supported by CUDA, or it errors immediately.

def Runtime.Autograd.Torch.Internal.EagerSession.add {α : Type} (s : EagerSession α) [Add α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise addition a + b. PyTorch: torch.add.

def Runtime.Autograd.Torch.Internal.EagerSession.sub {α : Type} (s : EagerSession α) [Sub α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise subtraction a - b. PyTorch: torch.sub.

def Runtime.Autograd.Torch.Internal.EagerSession.mul {α : Type} (s : EagerSession α) [Mul α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise multiplication a * b. PyTorch: torch.mul.

def Runtime.Autograd.Torch.Internal.EagerSession.scale {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Mul α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (c : α) : IO (TensorRef α sh)

Record scaling by a scalar constant. PyTorch: x * c.

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

Record elementwise absolute value. PyTorch: torch.abs.

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

Record elementwise square root. PyTorch: torch.sqrt.

def Runtime.Autograd.Torch.Internal.EagerSession.clamp {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (minVal maxVal : α) : IO (TensorRef α sh)

Record elementwise clamp to [minVal,maxVal]. PyTorch: torch.clamp.

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

Record elementwise maximum. PyTorch: torch.maximum.

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

Record elementwise minimum. PyTorch: torch.minimum.

def Runtime.Autograd.Torch.Internal.EagerSession.relu {α : Type} (s : EagerSession α) [Mul α] [Zero α] [Max α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise ReLU. PyTorch: torch.relu / torch.nn.functional.relu.

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

Record elementwise sigmoid. PyTorch: torch.sigmoid.

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

Record elementwise tanh. PyTorch: torch.tanh.

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

Record softmax (shape-preserving).

PyTorch comparison: torch.softmax(x, dim=...) (dimension convention is chosen by the underlying tape op).

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

Record stable log-softmax (shape-preserving, last-axis convention).

PyTorch comparison: torch.nn.functional.log_softmax(x, dim=-1).

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

Record elementwise softplus. PyTorch: torch.nn.functional.softplus.

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

Record elementwise exponential. PyTorch: torch.exp.

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

Record elementwise log. PyTorch: torch.log.

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

Record elementwise inverse 1/x. PyTorch: torch.reciprocal.

def Runtime.Autograd.Torch.Internal.EagerSession.safeLog {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (ε : α := Numbers.epsilon) : IO (TensorRef α sh)

Record elementwise log with epsilon guard.

PyTorch comparison: torch.log(torch.clamp(x, min=ε)).

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 α] [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))

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

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

Gather a scalar from a 1D vector with a Fin n index. PyTorch: x[i].

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRow {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : Fin rows) : IO (TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar))

Gather a row from a 2D tensor with a Fin rows index. PyTorch: x[i] for 2D tensors.

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

Gather a scalar from a 1D vector with a raw Nat index (totalized by the tape op).

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

Dynamic gather scalar using an index stored in NatRef.

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowRef {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : NatRef) : IO (TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar))

Dynamic gather row using an index stored in NatRef (out-of-range gives a zero row).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherVecNat {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n k : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim k Spec.Shape.scalar))

Gather k scalars using an explicit index tensor. PyTorch analogue: gather / advanced indexing.

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowsNat {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols k : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

Gather k rows using an explicit index tensor. PyTorch: index_select(dim=0, index=...).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherVecRef {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n k : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (idx : NatVecRef k) : IO (TensorRef α (Spec.Shape.dim k Spec.Shape.scalar))

Gather k scalars using indices stored in the nat-environment (NatVecRef).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowsRef {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols k : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : NatVecRef k) : IO (TensorRef α (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

Gather k rows using indices stored in the nat-environment (NatVecRef).

def Runtime.Autograd.Torch.Internal.EagerSession.scatterAddVec {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (v : TensorRef α Spec.Shape.scalar) (i : Fin n) : IO (TensorRef α (Spec.Shape.dim n Spec.Shape.scalar))

Scatter-add into a vector: return a copy of x with x[i] += v.

def Runtime.Autograd.Torch.Internal.EagerSession.scatterAddRow {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (v : TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar)) (i : Fin rows) : IO (TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)))

Scatter-add into a matrix row: return a copy of x with x[i,:] += v.

def Runtime.Autograd.Torch.Internal.EagerSession.linear {α : Type} (s : EagerSession α) [Inhabited α] [Add α] [Mul α] [Zero α] [DecidableEq Spec.Shape] [FastKernels.FastMatmul α] {inDim outDim : ℕ} (w : TensorRef α (Spec.Shape.dim outDim (Spec.Shape.dim inDim Spec.Shape.scalar))) (b : TensorRef α (Spec.Shape.dim outDim Spec.Shape.scalar)) (x : TensorRef α (Spec.Shape.dim inDim Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim outDim Spec.Shape.scalar))

Fully-connected linear layer y = w x + b (matvec).

If opts.fastKernels is enabled, uses a runtime-only fast kernel implementation. PyTorch comparison: torch.nn.functional.linear(x, weight=w, bias=b).

def Runtime.Autograd.Torch.Internal.EagerSession.mseLoss {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Inhabited α] [Add α] [Sub α] [Mul α] [Div α] [Zero α] [One α] [Coe ℕ α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (yhat target : TensorRef α sh) : IO (TensorRef α Spec.Shape.scalar)

Mean-squared-error loss returning a scalar.

If opts.fastKernels is enabled, uses a runtime-only fast kernel implementation. PyTorch comparison: torch.nn.functional.mse_loss(..., reduction=\"mean\").

def Runtime.Autograd.Torch.Internal.EagerSession.layerNorm {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x : TensorRef α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar))) (gamma beta : TensorRef α (Spec.Shape.dim embedDim Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)))

Layer normalization over embedding dimension. PyTorch: nn.LayerNorm / functional.layer_norm.

def Runtime.Autograd.Torch.Internal.EagerSession.batchnormChannelFirst {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {channels height width : ℕ} (h_c : channels > 0) (h_h : height > 0) (h_w : width > 0) (x : TensorRef α (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar)))) (gamma beta : TensorRef α (Spec.Shape.dim channels Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar))))

BatchNorm for channel-first images (C,H,W) (no batch axis). PyTorch: nn.BatchNorm2d (conceptually).

def Runtime.Autograd.Torch.Internal.EagerSession.multiHeadAttention {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : TensorRef α (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar))) (wo : TensorRef α (Spec.Shape.dim (numHeads * headDim) (Spec.Shape.dim dModel Spec.Shape.scalar))) (x : TensorRef α (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar))) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : IO (TensorRef α (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)))

Multi-head self-attention (typed, proof-friendly). PyTorch: nn.MultiheadAttention (conceptually).

def Runtime.Autograd.Torch.Internal.EagerSession.conv {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (w : TensorRef α (Spec.Shape.ofList (outC :: inC :: kernel.toList))) (b : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (x : TensorRef α (Spec.Shape.ofList (inC :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (outC :: (Spec.convOutSpatial inSpatial kernel stride padding).toList)))

N-D convolution for channels-first tensors (inC, spatial...) (no batch axis).

This is the generic counterpart to conv2d. PyTorch comparison: torch.nn.functional.conv{d}d specialized to a single sample.

def Runtime.Autograd.Torch.Internal.EagerSession.conv2d {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (input : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

2D convolution for channel-first images (C,H,W) (no batch axis). PyTorch: torch.nn.functional.conv2d.

def Runtime.Autograd.Torch.Internal.EagerSession.convTranspose {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (w : TensorRef α (Spec.Shape.ofList (inC :: outC :: kernel.toList))) (b : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (x : TensorRef α (Spec.Shape.ofList (inC :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (outC :: (Spec.convTransposeOutSpatial inSpatial kernel stride padding).toList)))

N-D transpose convolution for channels-first tensors (inC, spatial...) (no batch axis).

This is the generic counterpart to conv_transpose2d. PyTorch comparison: torch.nn.functional.conv_transpose{d}d specialized to a single sample.

def Runtime.Autograd.Torch.Internal.EagerSession.convTranspose2d {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim outC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (input : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim ((inH - 1) * stride - 2 * padding + kH) (Spec.Shape.dim ((inW - 1) * stride - 2 * padding + kW) Spec.Shape.scalar))))

2D transpose convolution for channel-first images (C,H,W) (no batch axis). PyTorch: torch.nn.functional.conv_transpose2d.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerSession.conv2dCompat {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (input : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for conv2d (compat shorthand).

def Runtime.Autograd.Torch.Internal.EagerSession.matmul {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {m n p : ℕ} (a : TensorRef α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : TensorRef α (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : IO (TensorRef α (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar)))

2D matrix multiplication. PyTorch: torch.matmul for 2D tensors.

def Runtime.Autograd.Torch.Internal.EagerSession.bmm {α : Type} (s : EagerSession α) [Add α] [Mul α] [Zero α] [DecidableEq Spec.Shape] {batch m n p : ℕ} (a : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)))) (b : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))))

Batched matrix multiplication. PyTorch: torch.bmm.

def Runtime.Autograd.Torch.Internal.EagerSession.concatVectors {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {n m : ℕ} (a : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (b : TensorRef α (Spec.Shape.dim m Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim (n + m) Spec.Shape.scalar))

Concatenate two vectors along dim 0. PyTorch: torch.cat([a,b], dim=0).

def Runtime.Autograd.Torch.Internal.EagerSession.concatDim0 {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {n m : ℕ} {sh : Spec.Shape} (a : TensorRef α (Spec.Shape.dim n sh)) (b : TensorRef α (Spec.Shape.dim m sh)) : IO (TensorRef α (Spec.Shape.dim (n + m) sh))

Concatenate along dim 0 for tensors with leading dimension. PyTorch: torch.cat(..., dim=0).

def Runtime.Autograd.Torch.Internal.EagerSession.sliceRange0 {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} {sh : Spec.Shape} (x : TensorRef α (Spec.Shape.dim n sh)) (start len : ℕ) (h : len + start ≤ n) : IO (TensorRef α (Spec.Shape.dim len sh))

Slice along dim 0: x[start:start+len]. PyTorch: standard slicing.

def Runtime.Autograd.Torch.Internal.EagerSession.maxPool {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D max pooling for channels-first tensors (C, spatial...) (no batch axis).

PyTorch comparison: torch.nn.functional.max_pool1d / max_pool2d / max_pool3d depending on the spatial rank d.

def Runtime.Autograd.Torch.Internal.EagerSession.avgPool {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D average pooling for channels-first tensors (C, spatial...) (no batch axis).

PyTorch comparison: torch.nn.functional.avg_pool1d / avg_pool2d / avg_pool3d depending on the spatial rank d.

def Runtime.Autograd.Torch.Internal.EagerSession.smoothMaxPool {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) (beta : α) : IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D smooth max pooling (log-sum-exp surrogate) for channels-first tensors (C, spatial...).

This is a differentiable approximation to max pooling; PyTorch does not expose it as a single primitive, but it can be emulated with logsumexp over local windows.

def Runtime.Autograd.Torch.Internal.EagerSession.maxPool2d {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

2D max-pooling (no batch axis). PyTorch: torch.nn.functional.max_pool2d.

def Runtime.Autograd.Torch.Internal.EagerSession.maxPool2dPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

2D max-pooling with padding (no batch axis). PyTorch: max_pool2d(..., padding=...).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerSession.maxPoolPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for max_pool2d_pad (PyTorch-style shorthand).

def Runtime.Autograd.Torch.Internal.EagerSession.smoothMaxPool2d {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) (beta : α) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

Smooth max-pooling (softmax pooling). Not a standard PyTorch primitive; see Torch.LinkedSession.smooth_max_pool2d.

def Runtime.Autograd.Torch.Internal.EagerSession.avgPool2d {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

2D average-pooling (no batch axis). PyTorch: torch.nn.functional.avg_pool2d.

def Runtime.Autograd.Torch.Internal.EagerSession.avgPool2dPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

2D average-pooling with padding (no batch axis). PyTorch: avg_pool2d(..., padding=...).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerSession.avgPoolPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for avg_pool2d_pad (PyTorch-style shorthand).

def Runtime.Autograd.Torch.Internal.EagerSession.backwardDenseAllCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (out : TensorRef α sh) (seed : Spec.Tensor α sh) : IO (Array Cuda.AnyBuffer)

Run reverse-mode backprop on the CUDA tape, returning device gradients for all tape entries.

This is the CUDA analogue of backwardDenseAll, but it does not download gradients back to the host. This is primarily useful for implementing GPU-native optimizer steps.

def Runtime.Autograd.Torch.Internal.EagerSession.backwardScalarDenseAllCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [One α] [DecidableEq Spec.Shape] (loss : TensorRef α Spec.Shape.scalar) : IO (Array Cuda.AnyBuffer)

Convenience wrapper for scalar losses on CUDA: backward with seed 1 (device buffers).

def Runtime.Autograd.Torch.Internal.EagerSession.backwardDenseAll {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (out : TensorRef α sh) (seed : Spec.Tensor α sh) : IO (Array (AnyTensor α))

Run reverse-mode backprop and return a dense gradient array for all tape entries.

seed is the upstream gradient for out (like PyTorch's backward(gradient=...)).

def Runtime.Autograd.Torch.Internal.EagerSession.backwardScalarDenseAll {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [One α] [DecidableEq Spec.Shape] (loss : TensorRef α Spec.Shape.scalar) : IO (Array (AnyTensor α))

Convenience wrapper for scalar losses: run backward with seed 1.

PyTorch comparison: loss.backward() for a scalar loss.

def Runtime.Autograd.Torch.Internal.EagerSession.grad {α : Type} {sh : Spec.Shape} [DecidableEq Spec.Shape] (grads : Array (AnyTensor α)) (x : TensorRef α sh) : IO (Spec.Tensor α sh)

Extract the gradient for a particular TensorRef from a dense gradient array.

def Runtime.Autograd.Torch.Internal.EagerSession.sgdStepAll {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Sub α] [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] (lr : α) (grads : Array (AnyTensor α)) : IO Unit

Apply an SGD update to all parameters recorded via use.

PyTorch comparison: for p in params: p.data -= lr * p.grad.

def Runtime.Autograd.Torch.Internal.EagerSession.sgdStepAllCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [DecidableEq Spec.Shape] (lr : α) (grads : Array Cuda.AnyBuffer) : IO Unit

Apply an SGD update to all parameters recorded via use, using CUDA device gradients.

This avoids downloading the full dense gradient array and keeps updated parameters in each Param's CUDA mirror. Host tensors are synchronized later by explicit parameter readback.

structure Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamParamState : Type

Device-side Adam moment buffers for one parameter leaf.

• m : Cuda.Buffer

  First moment buffer.

• v : Cuda.Buffer

  Second moment buffer.

• t : ℕ

  Adam step counter for this parameter.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamState : Type

def Runtime.Autograd.Torch.Internal.EagerSession.adamStepAllCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [DecidableEq Spec.Shape] (stateRef : IO.Ref CudaAdamState) (lr beta1 beta2 epsilon : α) (grads : Array Cuda.AnyBuffer) : IO Unit

Apply an Adam update to all parameters recorded via use, using CUDA device gradients.

This is the CUDA analogue of the generic TorchLean.Optim.adam path. It keeps Adam moments as device buffers and keeps updated parameters in each Param's CUDA mirror, so the next CUDA forward can reuse them without a host upload. Host tensors are synchronized later by explicit readback.

def Runtime.Autograd.Torch.Internal.EagerSession.adamWStepAllCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [DecidableEq Spec.Shape] (stateRef : IO.Ref CudaAdamState) (lr weightDecay beta1 beta2 epsilon : α) (grads : Array Cuda.AnyBuffer) : IO Unit

Apply an AdamW update to all parameters recorded via use, using CUDA device gradients.

This mirrors Optim.AdamW.update: moments are formed from the raw gradient, weight decay is applied directly to parameters, then the Adam update is applied. Like adamStepAllCuda, it keeps updated parameter buffers resident on device and only synchronizes the host copy when readback is requested.

Imperative sessions live in:

• Runtime.Autograd.TorchLean.Session (unified eager/compiled, recommended),
• Runtime.Autograd.Torch.Internal.SessionIR (proof-linked imperative session, internal).

torch.compile-style wrapper (cached static graph).

This is a thin wrapper around the proof-compiled graph model (GraphData) and its proven-correct reverse-mode accumulator (GraphData.backpropCtx). It compiles once, then you can call forward/backward repeatedly with new inputs.

Note: this does not cache a Runtime.Autograd.Tape for reuse across different inputs. The current tape compiler bakes the forward context into backward closures, so reusing a single tape across changing inputs would be unsound without redesigning the runtime node API.

structure Runtime.Autograd.Torch.CompiledScalar (α : Type) (Γ : List Spec.Shape) : Type

torch.compile-style wrapper for a scalar-valued computation over leaf context Γ.

This stores a proved node (NodeData) together with the preceding graph prefix so it can be evaluated and differentiated without rebuilding the whole graph each time.

• ssPrev : List Spec.Shape

  Shapes of internal SSA nodes preceding the scalar output node.

• gPrev : Proofs.Autograd.Algebra.GraphData α Unit Γ self.ssPrev

  Proved graph prefix that computes all preceding SSA nodes.

• node : Proofs.Autograd.Algebra.NodeData α Unit (Γ ++ self.ssPrev) Spec.Shape.scalar

  Final scalar output node over the leaf context plus graph prefix.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.CompiledScalar.TList (α : Type) (ss : List Spec.Shape) : Type

Convenience alias for the proved heterogeneous tensor list over a shape context.

def Runtime.Autograd.Torch.CompiledScalar.forward {α : Type} {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x : TList α Γ) : Spec.Tensor α Spec.Shape.scalar

Evaluate the scalar output for leaf values x.

def Runtime.Autograd.Torch.CompiledScalar.jvp {α : Type} {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x dx : TList α Γ) : Spec.Tensor α Spec.Shape.scalar

Forward-mode Jacobian-vector product (JVP) at x with tangent dx.

def Runtime.Autograd.Torch.CompiledScalar.backward {α : Type} [Add α] [Zero α] [One α] {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x : TList α Γ) : TList α Γ

Reverse-mode backprop for a scalar output with implicit seed 1.

Returns a TList of gradients aligned with the leaf context Γ.

def Runtime.Autograd.Torch.CompiledScalar.backwardWithSeed {α : Type} [Add α] [Zero α] {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x : TList α Γ) (seedOut : α) : TList α Γ

Reverse-mode backprop for a scalar output with an explicit scalar seed.

PyTorch comparison: loss.backward(gradient=seedOut) for a scalar loss.

torch.compile-style wrapper for tensor-valued outputs.

This is the same idea as CompiledScalar, but parameterized by an arbitrary output shape τ. It supports:

• forward (evaluate the output),
• jvp (forward-mode JVP, provided all ops supply jvp),
• vjpWithSeed (reverse-mode VJP with an explicit cotangent seed at the output).

structure Runtime.Autograd.Torch.CompiledOut (α : Type) (Γ : List Spec.Shape) (τ : Spec.Shape) : Type

torch.compile-style wrapper for a tensor-valued output of shape τ.

This generalizes CompiledScalar to arbitrary output shapes and provides forward-mode JVP and reverse-mode VJP (with explicit seed).

• ssPrev : List Spec.Shape

  Shapes of internal SSA nodes preceding the output node.

• gPrev : Proofs.Autograd.Algebra.GraphData α Unit Γ self.ssPrev

  Proved graph prefix that computes all preceding SSA nodes.

• node : Proofs.Autograd.Algebra.NodeData α Unit (Γ ++ self.ssPrev) τ

  Final output node over the leaf context plus graph prefix.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.CompiledOut.TList (α : Type) (ss : List Spec.Shape) : Type

Convenience alias for the proved heterogeneous tensor list over a shape context.

def Runtime.Autograd.Torch.CompiledOut.forward {α : Type} {Γ : List Spec.Shape} {τ : Spec.Shape} (c : CompiledOut α Γ τ) (x : TList α Γ) : Spec.Tensor α τ

Evaluate the output tensor for leaf values x.

def Runtime.Autograd.Torch.CompiledOut.jvp {α : Type} {Γ : List Spec.Shape} {τ : Spec.Shape} (c : CompiledOut α Γ τ) (x dx : TList α Γ) : Spec.Tensor α τ

Forward-mode Jacobian-vector product (JVP) at x with tangent dx.

def Runtime.Autograd.Torch.CompiledOut.vjpWithSeed {α : Type} [Add α] [Zero α] {Γ : List Spec.Shape} {τ : Spec.Shape} (c : CompiledOut α Γ τ) (x : TList α Γ) (seedOut : Spec.Tensor α τ) : TList α Γ

Reverse-mode vector-Jacobian product (VJP) with an explicit output cotangent seed.

This is the tensor-valued analogue of CompiledScalar.backwardWithSeed. PyTorch comparison: out.backward(gradient=seedOut) (for a tensor output).

def Runtime.Autograd.Torch.compileScalar {α : Type} [DecidableEq Spec.Shape] {Γ : List Spec.Shape} (build : Compiled.GraphM.M α Γ (Compiled.GraphM.Var Spec.Shape.scalar)) : Result (CompiledScalar α Γ)

Compile a scalar-output graph builder into a CompiledScalar.

The builder is expressed in the Compiled.GraphM monad. We expect it to produce at least one node and return a variable of scalar shape.

def Runtime.Autograd.Torch.compileOut {α : Type} [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {τ : Spec.Shape} (build : Compiled.GraphM.M α Γ (Compiled.GraphM.Var τ)) : Result (CompiledOut α Γ τ)

Compile a tensor-output graph builder into a CompiledOut.

We require that the returned Var τ is the last node produced by the builder, so the wrapper can store the prefix graph and final output node cleanly.

def Runtime.Autograd.Torch.Proofs.Autograd.Algebra.TList.append {α : Type} {ss₁ ss₂ : List Spec.Shape} : TList α ss₁ → TList α ss₂ → TList α (ss₁ ++ ss₂)

Append two TLists.

This is a small utility for bridging between curried APIs and list-of-shapes APIs.

def Runtime.Autograd.Torch.Proofs.Autograd.Algebra.TList.splitAppend {α : Type} {ss₁ ss₂ : List Spec.Shape} : TList α (ss₁ ++ ss₂) → TList α ss₁ × TList α ss₂

Split a TList α (ss₁ ++ ss₂) into its left and right parts.

This is the inverse of TList.append.

def Runtime.Autograd.Torch.Curried.Fn (α : Type) : List Spec.Shape → Type → Type

Type of a curried function accepting one tensor argument per shape in ss.

For example, Fn α [s₁, s₂] β is Tensor α s₁ → Tensor α s₂ → β.

def Runtime.Autograd.Torch.Curried.curry {α β : Type} {ss : List Spec.Shape} : (TList α ss → β) → Fn α ss β

Convert a function on TList inputs into its curried form.

def Runtime.Autograd.Torch.Curried.uncurry {α β : Type} {ss : List Spec.Shape} : Fn α ss β → TList α ss → β

Convert a curried function into a function on TList inputs.

RefList is the reference-analogue of TList: a heterogeneous list of Ref s values indexed by a shape list.

This is used to write backend-generic code over references (e.g. TensorRefs in eager mode, or GraphM.Vars in compiled mode).

inductive Runtime.Autograd.Torch.RefList (Ref : Spec.Shape → Type) : List Spec.Shape → Type

Reference-analogue of TList: a heterogeneous list of Ref s values indexed by shapes.

• nil {Ref : Spec.Shape → Type} : RefList Ref []
• cons {Ref : Spec.Shape → Type} {s : Spec.Shape} {ss : List Spec.Shape} : Ref s → RefList Ref ss → RefList Ref (s :: ss)

def Runtime.Autograd.Torch.RefList.append {Ref : Spec.Shape → Type} {ss₁ ss₂ : List Spec.Shape} : RefList Ref ss₁ → RefList Ref ss₂ → RefList Ref (ss₁ ++ ss₂)

Append two RefLists.

def Runtime.Autograd.Torch.RefList.split {Ref : Spec.Shape → Type} {ss₁ ss₂ : List Spec.Shape} : RefList Ref (ss₁ ++ ss₂) → RefList Ref ss₁ × RefList Ref ss₂

Split a RefList Ref (ss₁ ++ ss₂) into its left and right parts.

def Runtime.Autograd.Torch.RefList.splitAppend1 {Ref : Spec.Shape → Type} {ss : List Spec.Shape} {τ : Spec.Shape} : RefList Ref (ss ++ [τ]) → RefList Ref ss × Ref τ

Split a RefList Ref (ss ++ [τ]) into its prefix and last element.

def Runtime.Autograd.Torch.CurriedRef (Ref : Spec.Shape → Type) : List Spec.Shape → Type → Type

Type of a curried function over references, one Ref s argument per shape in ss.

This mirrors Curried.Fn, but for Ref-valued arguments (e.g. TensorRefs in eager mode or GraphM.Vars in compiled mode).

def Runtime.Autograd.Torch.CurriedRef.uncurry {Ref : Spec.Shape → Type} {β : Type} {ss : List Spec.Shape} : CurriedRef Ref ss β → RefList Ref ss → β

Uncurry a curried reference function to accept a RefList.

def Runtime.Autograd.Torch.CurriedRef.curry {Ref : Spec.Shape → Type} {β : Type} {ss : List Spec.Shape} : (RefList Ref ss → β) → CurriedRef Ref ss β

Curry a reference function that consumes a RefList.

def Runtime.Autograd.Torch.CurriedRef.applyVarList {Γ : List Spec.Shape} {β : Type} : CurriedRef (fun (s : Spec.Shape) => Compiled.GraphM.Var s) Γ β → Compiled.GraphM.VarList Γ → β

Apply a curried reference function to a GraphM.VarList.

This is a convenience for the compiled backend, where leaves/inputs are represented as Vars.

class Runtime.Autograd.Torch.Ops (m : Type → Type) (α : Type) [Context α] [DecidableEq Spec.Shape] : Type 1

Backend-generic interface for building and executing tensor programs.

This typeclass lets you write a single model/loss once (polymorphic over Ops m α) and then choose:

• an eager backend that executes immediately on a runtime tape, or
• a compiled backend that records proved IR (GraphM) for later compilation/proofs.

Each method corresponds to a Tensor op; implementations are expected to match the semantics of the corresponding Runtime.Autograd.Tape.* / Compiled.GraphM.* operator.

• Ref : Spec.Shape → Type
• const {s : Spec.Shape} : Spec.Tensor α s → m (Ref m α s)
• add {s : Spec.Shape} : Ref m α s → Ref m α s → m (Ref m α s)
• sub {s : Spec.Shape} : Ref m α s → Ref m α s → m (Ref m α s)
• mul {s : Spec.Shape} : Ref m α s → Ref m α s → m (Ref m α s)
• scale {s : Spec.Shape} : Ref m α s → α → m (Ref m α s)
• abs {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• sqrt {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• clamp {s : Spec.Shape} : Ref m α s → α → α → m (Ref m α s)
• max {s : Spec.Shape} : Ref m α s → Ref m α s → m (Ref m α s)
• min {s : Spec.Shape} : Ref m α s → Ref m α s → m (Ref m α s)
• broadcastTo {s₁ s₂ : Spec.Shape} : s₁.CanBroadcastTo s₂ → Ref m α s₁ → m (Ref m α s₂)
• reshape {s₁ s₂ : Spec.Shape} : Ref m α s₁ → (h : s₁.size = s₂.size) → m (Ref m α s₂)
• transpose2d {mDim nDim : ℕ} : Ref m α (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar)) → m (Ref m α (Spec.Shape.dim nDim (Spec.Shape.dim mDim Spec.Shape.scalar)))
• transpose3dFirstToLast {a b c : ℕ} : Ref m α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim b (Spec.Shape.dim c (Spec.Shape.dim a Spec.Shape.scalar))))
• transpose3dLastToFirst {a b c : ℕ} : Ref m α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim c (Spec.Shape.dim a (Spec.Shape.dim b Spec.Shape.scalar))))
• transpose3dLastTwo {a b c : ℕ} : Ref m α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim a (Spec.Shape.dim c (Spec.Shape.dim b Spec.Shape.scalar))))
• swapAdjacentAtDepth {s : Spec.Shape} (depth : ℕ) : Ref m α s → m (Ref m α (s.swapAdjacentAtDepth depth))
• reduceSum {s : Spec.Shape} (axis : ℕ) [Spec.Shape.valid_axis_inst axis s] [s.WellFormed] : Ref m α s → m (Ref m α (Spec.Tensor.shapeAfterSum s axis))
• reduceMean {s : Spec.Shape} (axis : ℕ) [Spec.Shape.valid_axis_inst axis s] [s.WellFormed] : Ref m α s → m (Ref m α (Spec.Tensor.shapeAfterSum s axis))
• gatherScalar {n : ℕ} : Ref m α (Spec.Shape.dim n Spec.Shape.scalar) → Fin n → m (Ref m α Spec.Shape.scalar)
• gatherRow {rows cols : ℕ} : Ref m α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)) → Fin rows → m (Ref m α (Spec.Shape.dim cols Spec.Shape.scalar))
• gatherScalarNat {n : ℕ} : Ref m α (Spec.Shape.dim n Spec.Shape.scalar) → ℕ → m (Ref m α Spec.Shape.scalar)
• gatherVecNat {n k : ℕ} : Ref m α (Spec.Shape.dim n Spec.Shape.scalar) → Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar) → m (Ref m α (Spec.Shape.dim k Spec.Shape.scalar))
• gatherRowsNat {rows cols k : ℕ} : Ref m α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)) → Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar) → m (Ref m α (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))
• scatterAddVec {n : ℕ} : Ref m α (Spec.Shape.dim n Spec.Shape.scalar) → Ref m α Spec.Shape.scalar → Fin n → m (Ref m α (Spec.Shape.dim n Spec.Shape.scalar))
• scatterAddRow {rows cols : ℕ} : Ref m α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim cols Spec.Shape.scalar) → Fin rows → m (Ref m α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)))
• matmul {mDim nDim pDim : ℕ} : Ref m α (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim nDim (Spec.Shape.dim pDim Spec.Shape.scalar)) → m (Ref m α (Spec.Shape.dim mDim (Spec.Shape.dim pDim Spec.Shape.scalar)))
• bmm {batch mDim nDim pDim : ℕ} : Ref m α (Spec.Shape.dim batch (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar))) → Ref m α (Spec.Shape.dim batch (Spec.Shape.dim nDim (Spec.Shape.dim pDim Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim batch (Spec.Shape.dim mDim (Spec.Shape.dim pDim Spec.Shape.scalar))))
• concatVectors {nDim mDim : ℕ} : Ref m α (Spec.Shape.dim nDim Spec.Shape.scalar) → Ref m α (Spec.Shape.dim mDim Spec.Shape.scalar) → m (Ref m α (Spec.Shape.dim (nDim + mDim) Spec.Shape.scalar))
• concatDim0 {nDim mDim : ℕ} {s : Spec.Shape} : Ref m α (Spec.Shape.dim nDim s) → Ref m α (Spec.Shape.dim mDim s) → m (Ref m α (Spec.Shape.dim (nDim + mDim) s))
• sliceRange0 {nDim : ℕ} {s : Spec.Shape} (start len : ℕ) (h : len + start ≤ nDim) : Ref m α (Spec.Shape.dim nDim s) → m (Ref m α (Spec.Shape.dim len s))
• maxPool {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} : Ref m α (Spec.Shape.ofList (C :: inSpatial.toList)) → m (Ref m α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))
• avgPool {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) : Ref m α (Spec.Shape.ofList (C :: inSpatial.toList)) → m (Ref m α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))
• smoothMaxPool {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} : Ref m α (Spec.Shape.ofList (C :: inSpatial.toList)) → α → m (Ref m α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))
• maxPool2d {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} : Ref m α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))
• maxPool2dPad {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} : Ref m α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))
• smoothMaxPool2d {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} : Ref m α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) → α → m (Ref m α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))
• avgPool2d {kH kW inH inW inC stride : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) : Ref m α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))
• avgPool2dPad {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) : Ref m α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))
• relu {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• sigmoid {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• tanh {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• softmax {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• logSoftmax {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• softplus {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• exp {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• log {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• inv {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• detach {s : Spec.Shape} : Ref m α s → m (Ref m α s)
• safeLog {s : Spec.Shape} : Ref m α s → α → m (Ref m α s)
• sum {s : Spec.Shape} : Ref m α s → m (Ref m α Spec.Shape.scalar)
• flatten {s : Spec.Shape} : Ref m α s → m (Ref m α (Spec.Shape.dim s.size Spec.Shape.scalar))
• linear {inDim outDim : ℕ} : Ref m α (Spec.Shape.dim outDim (Spec.Shape.dim inDim Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim outDim Spec.Shape.scalar) → Ref m α (Spec.Shape.dim inDim Spec.Shape.scalar) → m (Ref m α (Spec.Shape.dim outDim Spec.Shape.scalar))
• mseLoss {s : Spec.Shape} : Ref m α s → Ref m α s → m (Ref m α Spec.Shape.scalar)
• layerNorm {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) : Ref m α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim embedDim Spec.Shape.scalar) → Ref m α (Spec.Shape.dim embedDim Spec.Shape.scalar) → m (Ref m α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)))
• batchnormChannelFirst {channels height width : ℕ} (h_c : channels > 0) (h_h : height > 0) (h_w : width > 0) : Ref m α (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar))) → Ref m α (Spec.Shape.dim channels Spec.Shape.scalar) → Ref m α (Spec.Shape.dim channels Spec.Shape.scalar) → m (Ref m α (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar))))
• multiHeadAttention {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) : Ref m α (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim (numHeads * headDim) (Spec.Shape.dim dModel Spec.Shape.scalar)) → Ref m α (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)) → Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)))
• conv {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} : Ref m α (Spec.Shape.ofList (outC :: inC :: kernel.toList)) → Ref m α (Spec.Shape.dim outC Spec.Shape.scalar) → Ref m α (Spec.Shape.ofList (inC :: inSpatial.toList)) → m (Ref m α (Spec.Shape.ofList (outC :: (Spec.convOutSpatial inSpatial kernel stride padding).toList)))
• convTranspose {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} : Ref m α (Spec.Shape.ofList (inC :: outC :: kernel.toList)) → Ref m α (Spec.Shape.dim outC Spec.Shape.scalar) → Ref m α (Spec.Shape.ofList (inC :: inSpatial.toList)) → m (Ref m α (Spec.Shape.ofList (outC :: (Spec.convTransposeOutSpatial inSpatial kernel stride padding).toList)))
• conv2d {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} : Ref m α (Spec.Shape.dim outC (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar)))) → Ref m α (Spec.Shape.dim outC Spec.Shape.scalar) → Ref m α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim outC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))
• convTranspose2d {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} : Ref m α (Spec.Shape.dim inC (Spec.Shape.dim outC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar)))) → Ref m α (Spec.Shape.dim outC Spec.Shape.scalar) → Ref m α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) → m (Ref m α (Spec.Shape.dim outC (Spec.Shape.dim ((inH - 1) * stride - 2 * padding + kH) (Spec.Shape.dim ((inW - 1) * stride - 2 * padding + kW) Spec.Shape.scalar))))
• randUniform {s : Spec.Shape} (seed : ℕ) : m (Ref m α s)
• bernoulliMask {s : Spec.Shape} : Ref m α Spec.Shape.scalar → (seed : ℕ) → m (Ref m α s)

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Ref {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] (s : Spec.Shape) : Type

Reference type for the current Ops instance.

In eager mode this will typically be TensorRef; in compiled mode it will typically be GraphM.Var.

def Runtime.Autograd.Torch.const {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (t : Spec.Tensor α s) : m (Ref s)

Re-export of Ops.const. PyTorch: torch.tensor(...) / literal constants.

def Runtime.Autograd.Torch.add {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (a b : Ref s) : m (Ref s)

Re-export of Ops.add. PyTorch: torch.add / +.

def Runtime.Autograd.Torch.sub {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (a b : Ref s) : m (Ref s)

Re-export of Ops.sub. PyTorch: torch.sub / -.

def Runtime.Autograd.Torch.mul {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (a b : Ref s) : m (Ref s)

Re-export of Ops.mul. PyTorch: torch.mul / *.

def Runtime.Autograd.Torch.scale {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) (c : α) : m (Ref s)

Re-export of Ops.scale. PyTorch: x * c for a scalar c.

def Runtime.Autograd.Torch.abs {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.abs. PyTorch: torch.abs.

def Runtime.Autograd.Torch.sqrt {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.sqrt. PyTorch: torch.sqrt.

def Runtime.Autograd.Torch.clamp {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) (minVal maxVal : α) : m (Ref s)

Re-export of Ops.clamp. PyTorch: torch.clamp.

def Runtime.Autograd.Torch.max {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (a b : Ref s) : m (Ref s)

Re-export of Ops.max. PyTorch: torch.maximum.

def Runtime.Autograd.Torch.min {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (a b : Ref s) : m (Ref s)

Re-export of Ops.min. PyTorch: torch.minimum.

def Runtime.Autograd.Torch.broadcastTo {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s₁ s₂ : Spec.Shape} (cb : s₁.CanBroadcastTo s₂) (x : Ref s₁) : m (Ref s₂)

Re-export of Ops.broadcastTo. PyTorch: broadcasting / expand.

def Runtime.Autograd.Torch.reshape {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s₁ s₂ : Spec.Shape} (x : Ref s₁) (h : s₁.size = s₂.size) : m (Ref s₂)

Re-export of Ops.reshape. PyTorch: reshape / view.

def Runtime.Autograd.Torch.transpose2d {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {mDim nDim : ℕ} (x : Ref (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar))) : m (Ref (Spec.Shape.dim nDim (Spec.Shape.dim mDim Spec.Shape.scalar)))

Re-export of Ops.transpose2d. PyTorch: x.t() / transpose.

def Runtime.Autograd.Torch.transpose3dFirstToLast {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {a b c : ℕ} (x : Ref (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim b (Spec.Shape.dim c (Spec.Shape.dim a Spec.Shape.scalar))))

Re-export of Ops.transpose3d_first_to_last. PyTorch: permute(1,2,0).

def Runtime.Autograd.Torch.transpose3dLastToFirst {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {a b c : ℕ} (x : Ref (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim c (Spec.Shape.dim a (Spec.Shape.dim b Spec.Shape.scalar))))

Re-export of Ops.transpose3d_last_to_first. PyTorch: permute(2,0,1).

def Runtime.Autograd.Torch.transpose3dLastTwo {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {a b c : ℕ} (x : Ref (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim a (Spec.Shape.dim c (Spec.Shape.dim b Spec.Shape.scalar))))

Re-export of Ops.transpose3d_last_two. PyTorch: transpose(1,2).

def Runtime.Autograd.Torch.swapAdjacentAtDepth {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (depth : ℕ) (x : Ref s) : m (Ref (s.swapAdjacentAtDepth depth))

Re-export of Ops.swapAdjacentAtDepth (general adjacent-axis swap).

def Runtime.Autograd.Torch.reduceSum {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.valid_axis_inst axis s] [s.WellFormed] (x : Ref s) : m (Ref (Spec.Tensor.shapeAfterSum s axis))

Re-export of Ops.reduce_sum. PyTorch: torch.sum(..., dim=axis).

def Runtime.Autograd.Torch.reduceMean {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.valid_axis_inst axis s] [s.WellFormed] (x : Ref s) : m (Ref (Spec.Tensor.shapeAfterSum s axis))

Re-export of Ops.reduce_mean. PyTorch: torch.mean(..., dim=axis).

def Runtime.Autograd.Torch.gatherScalar {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {n : ℕ} (x : Ref (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin n) : m (Ref Spec.Shape.scalar)

Re-export of Ops.gather_scalar. PyTorch: x[i] (1D).

def Runtime.Autograd.Torch.gatherRow {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {rows cols : ℕ} (x : Ref (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : Fin rows) : m (Ref (Spec.Shape.dim cols Spec.Shape.scalar))

Re-export of Ops.gather_row. PyTorch: x[i] (2D row).

def Runtime.Autograd.Torch.gatherScalarNat {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {n : ℕ} (x : Ref (Spec.Shape.dim n Spec.Shape.scalar)) (i : ℕ) : m (Ref Spec.Shape.scalar)

Re-export of Ops.gather_scalar_nat (index is a raw Nat).

def Runtime.Autograd.Torch.gatherVecNat {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {n k : ℕ} (x : Ref (Spec.Shape.dim n Spec.Shape.scalar)) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : m (Ref (Spec.Shape.dim k Spec.Shape.scalar))

Re-export of Ops.gather_vec_nat (index tensor).

def Runtime.Autograd.Torch.gatherRowsNat {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {rows cols k : ℕ} (x : Ref (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : m (Ref (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

Re-export of Ops.gather_rows_nat (index tensor).

def Runtime.Autograd.Torch.scatterAddVec {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {n : ℕ} (x : Ref (Spec.Shape.dim n Spec.Shape.scalar)) (v : Ref Spec.Shape.scalar) (i : Fin n) : m (Ref (Spec.Shape.dim n Spec.Shape.scalar))

Re-export of Ops.scatter_add_vec.

def Runtime.Autograd.Torch.scatterAddRow {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {rows cols : ℕ} (x : Ref (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (v : Ref (Spec.Shape.dim cols Spec.Shape.scalar)) (i : Fin rows) : m (Ref (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)))

Re-export of Ops.scatter_add_row.

def Runtime.Autograd.Torch.matmul {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {mDim nDim pDim : ℕ} (a : Ref (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar))) (b : Ref (Spec.Shape.dim nDim (Spec.Shape.dim pDim Spec.Shape.scalar))) : m (Ref (Spec.Shape.dim mDim (Spec.Shape.dim pDim Spec.Shape.scalar)))

Re-export of Ops.matmul. PyTorch: torch.matmul for 2D tensors.

def Runtime.Autograd.Torch.bmm {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {batch mDim nDim pDim : ℕ} (a : Ref (Spec.Shape.dim batch (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar)))) (b : Ref (Spec.Shape.dim batch (Spec.Shape.dim nDim (Spec.Shape.dim pDim Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim batch (Spec.Shape.dim mDim (Spec.Shape.dim pDim Spec.Shape.scalar))))

Re-export of Ops.bmm. PyTorch: torch.bmm.

def Runtime.Autograd.Torch.concatVectors {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {nDim mDim : ℕ} (a : Ref (Spec.Shape.dim nDim Spec.Shape.scalar)) (b : Ref (Spec.Shape.dim mDim Spec.Shape.scalar)) : m (Ref (Spec.Shape.dim (nDim + mDim) Spec.Shape.scalar))

Re-export of Ops.concat_vectors. PyTorch: torch.cat([a,b], dim=0) for vectors.

def Runtime.Autograd.Torch.concatDim0 {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {nDim mDim : ℕ} {s : Spec.Shape} (a : Ref (Spec.Shape.dim nDim s)) (b : Ref (Spec.Shape.dim mDim s)) : m (Ref (Spec.Shape.dim (nDim + mDim) s))

Re-export of Ops.concat_dim0. PyTorch: torch.cat(..., dim=0).

def Runtime.Autograd.Torch.sliceRange0 {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {nDim : ℕ} {s : Spec.Shape} (start len : ℕ) (h : len + start ≤ nDim) (x : Ref (Spec.Shape.dim nDim s)) : m (Ref (Spec.Shape.dim len s))

Re-export of Ops.slice_range0. PyTorch: x[start:start+len] on the leading dimension.

def Runtime.Autograd.Torch.maxPool {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : Ref (Spec.Shape.ofList (C :: inSpatial.toList))) : m (Ref (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

Re-export of Ops.max_pool (generic N-D max pooling, channels-first; no batch axis).

PyTorch comparison: torch.nn.functional.max_pool1d / max_pool2d / max_pool3d depending on the spatial rank d.

def Runtime.Autograd.Torch.avgPool {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (x : Ref (Spec.Shape.ofList (C :: inSpatial.toList))) : m (Ref (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

Re-export of Ops.avg_pool (generic N-D average pooling, channels-first; no batch axis).

PyTorch comparison: torch.nn.functional.avg_pool1d / avg_pool2d / avg_pool3d depending on the spatial rank d.

def Runtime.Autograd.Torch.smoothMaxPool {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : Ref (Spec.Shape.ofList (C :: inSpatial.toList))) (beta : α) : m (Ref (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

Re-export of Ops.smooth_max_pool (generic N-D smooth max pooling, channels-first; no batch axis).

This is a differentiable approximation to max pooling; PyTorch does not expose it as a single primitive.

def Runtime.Autograd.Torch.maxPool2d {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

Re-export of Ops.max_pool2d. PyTorch: torch.nn.functional.max_pool2d.

def Runtime.Autograd.Torch.maxPool2dPad {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Re-export of Ops.max_pool2d_pad. PyTorch: max_pool2d(..., padding=...).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.maxPoolPad {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for max_pool2d_pad (PyTorch-style shorthand).

def Runtime.Autograd.Torch.smoothMaxPool2d {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) (beta : α) : m (Ref (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

Re-export of Ops.smooth_max_pool2d (softmax pooling).

def Runtime.Autograd.Torch.avgPool2d {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {kH kW inH inW inC stride : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

Re-export of Ops.avg_pool2d. PyTorch: torch.nn.functional.avg_pool2d.

def Runtime.Autograd.Torch.avgPool2dPad {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Re-export of Ops.avg_pool2d_pad. PyTorch: avg_pool2d(..., padding=...).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.avgPoolPad {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for avg_pool2d_pad (PyTorch-style shorthand).

def Runtime.Autograd.Torch.relu {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.relu.

def Runtime.Autograd.Torch.sigmoid {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.sigmoid.

def Runtime.Autograd.Torch.tanh {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.tanh.

def Runtime.Autograd.Torch.softmax {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.softmax.

def Runtime.Autograd.Torch.softplus {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.softplus.

def Runtime.Autograd.Torch.exp {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.exp.

def Runtime.Autograd.Torch.log {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.log.

def Runtime.Autograd.Torch.inv {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.inv (reciprocal).

def Runtime.Autograd.Torch.detach {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

Re-export of Ops.detach. PyTorch: x.detach().

def Runtime.Autograd.Torch.safeLog {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) (ε : α := Numbers.epsilon) : m (Ref s)

Re-export of Ops.safe_log.

def Runtime.Autograd.Torch.randUniform {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (seed : ℕ) : m (Ref s)

Re-export of Ops.rand_uniform (deterministic seeded RNG).

def Runtime.Autograd.Torch.bernoulliMask {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (keepProb : Ref Spec.Shape.scalar) (seed : ℕ) : m (Ref s)

Re-export of Ops.bernoulli_mask (deterministic dropout-style mask).

def Runtime.Autograd.Torch.logSoftmax {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) (ε : α := Numbers.epsilon) : m (Ref s)

Stable log_softmax(x) along the last axis.

This is a backend primitive with the standard max-shifted formulation x - max(x) - log(sum(exp(x - max(x)))), matching PyTorch's numerical intent. The optional ε parameter is kept for source compatibility with existing TorchLean callers and is intentionally ignored; callers that need an epsilon-smoothed logarithm should use safeLog explicitly.

def Runtime.Autograd.Torch.silu {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Monad m] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

SiLU / swish: x * sigmoid(x).

def Runtime.Autograd.Torch.gelu {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Monad m] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref s)

GELU (approximation used by many ML frameworks):

0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x^3))).

This is defined using existing primitives (tanh, mul, add, scale), so it works in eager, compiled, and verifier-IR backends without introducing a new opcode.

def Runtime.Autograd.Torch.globalAvgPool2dChw {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Monad m] [Ops m α] {c h w : ℕ} (h_c_pos : c > 0) (h_h_pos : h > 0) (h_w_pos : w > 0) (x : Ref (Spec.Shape.dim c (Spec.Shape.dim h (Spec.Shape.dim w Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim c Spec.Shape.scalar))

Global average pooling over the last two axes of a C×H×W tensor (channel-first).

Returns a vector C, averaging each channel over H×W.

def Runtime.Autograd.Torch.globalAvgPool2dNchw {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Monad m] [Ops m α] {n c h w : ℕ} (h_n_pos : n > 0) (h_c_pos : c > 0) (h_h_pos : h > 0) (h_w_pos : w > 0) (x : Ref (Spec.Shape.dim n (Spec.Shape.dim c (Spec.Shape.dim h (Spec.Shape.dim w Spec.Shape.scalar))))) : m (Ref (Spec.Shape.dim n (Spec.Shape.dim c Spec.Shape.scalar)))

Global average pooling over the last two axes of an N×C×H×W tensor (PyTorch default layout).

Returns N×C, averaging each channel over H×W for each batch element.

def Runtime.Autograd.Torch.sum {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref Spec.Shape.scalar)

Re-export of Ops.sum. PyTorch: x.sum().

def Runtime.Autograd.Torch.flatten {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (x : Ref s) : m (Ref (Spec.Shape.dim s.size Spec.Shape.scalar))

Re-export of Ops.flatten. PyTorch: torch.flatten.

def Runtime.Autograd.Torch.linear {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {inDim outDim : ℕ} (w : Ref (Spec.Shape.dim outDim (Spec.Shape.dim inDim Spec.Shape.scalar))) (b : Ref (Spec.Shape.dim outDim Spec.Shape.scalar)) (x : Ref (Spec.Shape.dim inDim Spec.Shape.scalar)) : m (Ref (Spec.Shape.dim outDim Spec.Shape.scalar))

Re-export of Ops.linear. PyTorch: torch.nn.functional.linear.

def Runtime.Autograd.Torch.mseLoss {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {s : Spec.Shape} (yhat target : Ref s) : m (Ref Spec.Shape.scalar)

Re-export of Ops.mse_loss. PyTorch: torch.nn.functional.mse_loss.

def Runtime.Autograd.Torch.layerNorm {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x : Ref (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar))) (gamma beta : Ref (Spec.Shape.dim embedDim Spec.Shape.scalar)) : m (Ref (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)))

Re-export of Ops.layer_norm. PyTorch: nn.LayerNorm / functional.layer_norm.

def Runtime.Autograd.Torch.batchnormChannelFirst {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {channels height width : ℕ} (h_c : channels > 0) (h_h : height > 0) (h_w : width > 0) (x : Ref (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar)))) (gamma beta : Ref (Spec.Shape.dim channels Spec.Shape.scalar)) : m (Ref (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar))))

Re-export of Ops.batchnorm_channel_first. PyTorch: nn.BatchNorm2d (conceptually).

def Runtime.Autograd.Torch.multiHeadAttention {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : Ref (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar))) (wo : Ref (Spec.Shape.dim (numHeads * headDim) (Spec.Shape.dim dModel Spec.Shape.scalar))) (x : Ref (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar))) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : m (Ref (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)))

Re-export of Ops.multi_head_attention.

def Runtime.Autograd.Torch.conv {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (weight : Ref (Spec.Shape.ofList (outC :: inC :: kernel.toList))) (bias : Ref (Spec.Shape.dim outC Spec.Shape.scalar)) (input : Ref (Spec.Shape.ofList (inC :: inSpatial.toList))) : m (Ref (Spec.Shape.ofList (outC :: (Spec.convOutSpatial inSpatial kernel stride padding).toList)))

Re-export of Ops.conv (generic N-D convolution, channels-first).

PyTorch comparison: torch.nn.functional.conv{d}d specialized to a single sample (no batch axis).

def Runtime.Autograd.Torch.convTranspose {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (weight : Ref (Spec.Shape.ofList (inC :: outC :: kernel.toList))) (bias : Ref (Spec.Shape.dim outC Spec.Shape.scalar)) (input : Ref (Spec.Shape.ofList (inC :: inSpatial.toList))) : m (Ref (Spec.Shape.ofList (outC :: (Spec.convTransposeOutSpatial inSpatial kernel stride padding).toList)))

Re-export of Ops.conv_transpose (generic N-D transpose convolution, channels-first).

PyTorch comparison: torch.nn.functional.conv_transpose{d}d specialized to a single sample.

def Runtime.Autograd.Torch.conv2d {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : Ref (Spec.Shape.dim outC (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : Ref (Spec.Shape.dim outC Spec.Shape.scalar)) (input : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim outC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Re-export of Ops.conv2d. PyTorch: torch.nn.functional.conv2d (conceptually, no batch axis).

def Runtime.Autograd.Torch.convTranspose2d {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : Ref (Spec.Shape.dim inC (Spec.Shape.dim outC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : Ref (Spec.Shape.dim outC Spec.Shape.scalar)) (input : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim outC (Spec.Shape.dim ((inH - 1) * stride - 2 * padding + kH) (Spec.Shape.dim ((inW - 1) * stride - 2 * padding + kW) Spec.Shape.scalar))))

Re-export of Ops.conv_transpose2d. PyTorch: torch.nn.functional.conv_transpose2d.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.conv2dCompat {m : Type → Type} {α : Type} [Context α] [DecidableEq Spec.Shape] [Ops m α] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : Ref (Spec.Shape.dim outC (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : Ref (Spec.Shape.dim outC Spec.Shape.scalar)) (input : Ref (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : m (Ref (Spec.Shape.dim outC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for conv2d (compat shorthand).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerM (α : Type) : Type → Type

Monad used for the eager Ops instance: read an Internal.EagerSession α and execute in IO.

This is the backend that makes Ops programs execute immediately by mutating a hidden runtime tape.

@[implicit_reducible]

instance Runtime.Autograd.Torch.instOpsEagerMOfTensorConv {α : Type} [Context α] [Internal.CudaBridge.TensorConv α] [DecidableEq Spec.Shape] : Ops (Internal.EagerM α) α

Ops instance for the eager Torch-style runtime.

This interprets Ops primitives by immediately executing them against the hidden mutable tape in the current Internal.EagerSession.

@[implicit_reducible]

instance Runtime.Autograd.Torch.instOpsM {α : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} : Ops (Compiled.GraphM.M α Γ) α

Ops instance for the compiled graph-building monad GraphM.

This interprets Ops primitives by recording typed IR nodes (rather than executing immediately). See Runtime.Autograd.Compiled.GraphM and Torch.LinkedSession for how these graphs are later run.

inductive Runtime.Autograd.Torch.ParamList (α : Type) : List Spec.Shape → Type

Heterogeneous list of trainable parameters, indexed by a list of shapes.

This is the Torch front-end analogue of "a parameter vector" (like model.parameters() in PyTorch), but with shapes tracked at the type level.

• nil {α : Type} : ParamList α []
• cons {α : Type} {s : Spec.Shape} {ss : List Spec.Shape} : Param α s → ParamList α ss → ParamList α (s :: ss)

def Runtime.Autograd.Torch.ParamList.subScaleMaterialize {α : Type} [Sub α] [Mul α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s → α → Spec.Tensor α s

Materialize the SGD update v - lr * g in a single traversal.

This is used by sgdStep_fast as a runtime-performance optimization to avoid building deep thunk chains when training for many steps.

def Runtime.Autograd.Torch.ParamList.ofTList {α : Type} {ss : List Spec.Shape} (xs : TList α ss) : IO (ParamList α ss)

Allocate a fresh ParamList from an initial TList of parameter tensors.

Each tensor becomes an IO.Ref so it can be updated by optimizer steps.

def Runtime.Autograd.Torch.ParamList.ofTListWithRequiresGrad {α : Type} {ss : List Spec.Shape} : TList α ss → List Bool → IO (ParamList α ss)

Allocate a fresh ParamList from an initial TList of parameter tensors, with explicit requiresGrad flags.

Returns an error when the flag list length does not match the parameter shape list length.

def Runtime.Autograd.Torch.ParamList.values {α : Type} {ss : List Spec.Shape} : ParamList α ss → IO (TList α ss)

Read the current parameter values as a TList aligned with the shape list.

def Runtime.Autograd.Torch.ParamList.valuesSynced {α : Type} [Internal.CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {ss : List Spec.Shape} : ParamList α ss → IO (TList α ss)

Read parameter values, synchronizing CUDA-resident mirrors first when necessary.

def Runtime.Autograd.Torch.ParamList.setValues {α : Type} {ss : List Spec.Shape} : ParamList α ss → TList α ss → IO Unit

Overwrite the current parameter values from a TList aligned with the shape list.

def Runtime.Autograd.Torch.ParamList.sgdStep {α : Type} [Context α] {ss : List Spec.Shape} : ParamList α ss → (lr : α) → TList α ss → IO Unit

Apply an SGD step p := p - lr * g to each parameter that has requiresGrad = true.

gs must be aligned with the parameter shapes.

def Runtime.Autograd.Torch.ParamList.sgdStepFast {α : Type} [Context α] {ss : List Spec.Shape} : ParamList α ss → (lr : α) → TList α ss → IO Unit

Like sgdStep, but uses a fully materialized update (subScaleMaterialize) for speed.

This is a runtime performance knob; mathematically it is equivalent to sgdStep.

structure Runtime.Autograd.Torch.ScalarTrainer (α : Type) (paramShapes inputShapes : List Spec.Shape) : Type

Bundle a scalar-loss training loop for a fixed parameter pack and input signature.

This is intended for simple demos:

• forward computes a scalar loss,
• backward computes gradients w.r.t. parameters,
• step applies an optimizer update (typically SGD),
• getParams reads current parameter values.

• params : ParamList α paramShapes

  Mutable trainable parameter pack.

• forward : Curried.Fn α inputShapes (IO (Spec.Tensor α Spec.Shape.scalar))

  Compute the scalar loss for a curried input pack.

• backward : Curried.Fn α inputShapes (IO (TList α paramShapes))

  Compute gradients aligned with paramShapes for a curried input pack.

• step : α → Curried.Fn α inputShapes (IO Unit)

  Apply one SGD-style update for a curried input pack.

• adamStep? : Option (α → α → α → α → Curried.Fn α inputShapes (IO Unit))

  Optional Adam update path.

  In eager CUDA mode this is a device-gradient/device-moment update path. Other backends expose none and should use the generic optimizer wrappers.

• adamWStep? : Option (α → α → α → α → α → Curried.Fn α inputShapes (IO Unit))

  Optional AdamW update path.

  In eager CUDA mode this is a device-gradient/device-moment update path with decoupled weight decay. Other backends expose none and should use the generic optimizer wrappers.

• getParams : IO (TList α paramShapes)

  Read current parameter values, synchronizing device mirrors if needed.

def Runtime.Autograd.Torch.Internal.gradsOfRefs {α : Type} [DecidableEq Spec.Shape] {ss : List Spec.Shape} : Array (AnyTensor α) → RefList (TensorRef α) ss → IO (TList α ss)

Extract gradients (as a typed TList) for a list of eager TensorRefs from a dense gradient array.

def Runtime.Autograd.Torch.Internal.useParams {α : Type} [CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {ss : List Spec.Shape} : ParamList α ss → EagerM α (RefList (TensorRef α) ss)

Record all parameters as tape leaves in an eager session, returning their corresponding TensorRefs.

This is the eager analogue of "using" a parameter pack during a forward pass.

def Runtime.Autograd.Torch.Internal.useInputs {α : Type} [CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {ss : List Spec.Shape} : TList α ss → EagerM α (RefList (TensorRef α) ss)

Record all input tensors as tape leaves in an eager session, returning their corresponding TensorRefs.

def Runtime.Autograd.Torch.scalarTrainer {α : Type} [Context α] [Internal.CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (opts : Options := { }) (initRequiresGrad : List Bool := List.replicate paramShapes.length true) (loss : {m : Type → Type} → [Monad m] → [inst : Ops m α] → CurriedRef (fun (s : Spec.Shape) => Ops.Ref m α s) (paramShapes ++ inputShapes) (m (Ops.Ref m α Spec.Shape.scalar))) : Curried.Fn α paramShapes (IO (ScalarTrainer α paramShapes inputShapes))

Build a ScalarTrainer from an initial parameter pack and a backend-generic loss definition.

loss is written once against the Ops interface over a concatenated context paramShapes ++ inputShapes. Depending on opts.backend, we either:

• compile the loss once (compiled backend), or
• execute it eagerly by building a runtime tape each step (eager backend).

def Runtime.Autograd.Torch.scalarTrainer.gradsPrefix {α : Type} [DecidableEq Spec.Shape] {ss : List Spec.Shape} : Array (AnyTensor α) → ℕ → IO (TList α ss)