Custom operations on tensors

Inputs use the public Arguments container for shape-indexed tensors. The pure evaluator uses checked row-major reads and stops at the first failure. Native execution transfers complete scalar words; byte arrays remain an implementation detail of the LibTorch bridge.

def NN.Kernel.Reader.ofArguments {α : Type} [TorchLean.Storage α] {shapes : List Spec.Shape} (inputs : TorchLean.Arguments α shapes) : Reader α

Read tensor arguments by operand number and row-major unsigned index.

The shape determines the bounds, without materializing an intermediate array. Invalid operand numbers and invalid addresses retain their distinct error diagnostics.

theorem NN.Kernel.Reader.ofArguments_eq_native {α : Type} [TorchLean.Storage α] {shapes : List Spec.Shape} (inputs : TorchLean.Arguments α shapes) (operand : ℕ) (hOperand : operand < shapes.length) (size : UInt64) (hSize : size.toNat = (shapes.get ⟨operand, hOperand⟩).size) (load : UInt64 → α) (hLoad : ∀ (index : UInt64) (hIndex : index.toNat < (shapes.get ⟨operand, hOperand⟩).size), load index = TorchLean.Tensor.Internal.Rep.getFlat (inputs.get ⟨operand, hOperand⟩) ⟨index.toNat, ⋯⟩) (index : UInt64) : ofArguments inputs operand index = if index ≥ size then Except.error (Error.bounds operand index size.toNat) else Except.ok (load index)

Tensor argument reads use the emitted unsigned bounds condition when the native length agrees with the tensor shape and in-bounds native loads agree with its entries.

The premises identify the buffer metadata and contents; they do not assert that a foreign memory implementation satisfies them. Out-of-bounds loads need no value because that branch never loads.

def NN.Kernel.Program.eval {α : Type} [TorchLean.Storage α] {shapes : List Spec.Shape} (program : Program α) (inputs : TorchLean.Arguments α shapes) (shape : Spec.Shape) : Except Error (TorchLean.Tensor α shape)

Evaluate a custom operation once per output entry, retaining the requested tensor shape.

All output indices must fit the expression language's unsigned addressing. An empty output makes no reads. Input failures stop construction in increasing output-index order.

theorem NN.Kernel.Program.eval_eq_reference {α : Type} [TorchLean.Storage α] {shapes : List Spec.Shape} (program : Program α) (inputs : TorchLean.Arguments α shapes) (shape : Spec.Shape) : program.eval inputs shape = if shape.size ≤ UInt64.size then TorchLean.Tensor.generateFlatM shape fun (index : Fin shape.size) => program.reference (Reader.ofArguments inputs) (↑index).toUInt64 else Except.error (Error.size shape.size)

Tensor evaluation agrees with the original source function at every output index.

The equality includes read failures and output-size rejection, not just successful tensors.

def NN.Kernel.Program.run {α : Type} [TorchLean.Storage α] {shapes : List Spec.Shape} (program : Program α) (inputs : TorchLean.Arguments α shapes) (shape : Spec.Shape) (device : Backend.Device := Backend.Device.cpu) : IO (TorchLean.Tensor α shape)

Execute a custom operation on the selected device, retaining the requested output shape.

CPU uses the Lean evaluator; GPU uses generated CUDA. A format outside the GPU implementation's limits uses the same Lean evaluator with a diagnostic, without changing precision. Other devices and operational failures are rejected. Native execution remains an external trust boundary.