Ordinary functions on tensors

f.run input (device := cpu) applies an ordinary Lean scalar function to a tensor. f.run input (device := gpu) uses generated CUDA for supported scalar functions, and the existing LibTorch operations for supported whole-tensor functions. grad := true records either kind through the existing operation interface and returns its forward value and input pullback. f.zip left right combines two equally shaped tensors with a binary scalar function. Captured parameters are constants. Scalar mapping preserves the input shape; whole-tensor functions can reduce or change shape. The call site infers the element type, storage and result shape. Named numerical recurrences can use the same API when their recursive argument is last and their parameters stay fixed. GPU lowering derives a loop and checks agreement with their Lean equations. CPU execution does not require GPU-compatible arithmetic. A scalar type without a GPU encoding runs on CPU with a diagnostic, retaining its type and precision. Unsupported source for an otherwise supported GPU type, compilation failures and execution failures remain errors.

def NN.Kernel.Frontend.tacticPrepare_tensor_function : Lean.ParserDescr

Prepare a certified GPU representation of a unary or binary scalar function.

class Function.Internal.Signature (α ι ο : Type) [TorchLean.Storage α] (input : Spec.Shape) (output : outParam Spec.Shape) : Prop

Determine the result shape from the ordinary function's signature, before call-site recording. This lets indexing and other consumers infer their types without waiting for source elaboration.

@[instance 10000, defaultInstance 1000]

instance Function.Internal.instSignature {α : Type} [TorchLean.Storage α] {input : Spec.Shape} : Signature α α α input input

instance Function.Internal.instSignatureTensor {α : Type} [TorchLean.Storage α] {input output : Spec.Shape} : Signature α (TorchLean.Tensor α input) (TorchLean.Tensor α output) input output

instance Function.Internal.instSignatureTensorScalar {α : Type} [TorchLean.Storage α] {input : Spec.Shape} : Signature α (TorchLean.Tensor α input) α input Spec.Shape.scalar

structure Function.Internal.Execution {ι ο : Type} (f : ι → ο) (α : Type) [TorchLean.Storage α] (input output : Spec.Shape) (grad : Bool) : Type

Call-site execution selected from the function's input and output types.

• run : TorchLean.Tensor α input → NN.Backend.Device → IO (match grad with | false => TorchLean.Tensor α output | true => TorchLean.autograd.Result α input output)

def Function.Internal.scalar {α : Type} [TorchLean.Storage α] {shape : Spec.Shape} (f : α → α) (input : TorchLean.Tensor α shape) (device : NN.Backend.Device := NN.Backend.Device.cpu) (gpuScalar? : Option (NN.Kernel.Scalar α) := by first | exact some inferInstance | exact none) (compilation : Except String (NN.Kernel.Elementwise f) := by prepare_tensor_function) : IO (TorchLean.Tensor α shape)

Apply a scalar function to every tensor entry on the selected device, preserving its shape.

CPU evaluates the original Lean function. GPU supports native FP32/FP64 and configured binary arithmetic. Types without a GPU encoding run on CPU with a diagnostic, without narrowing values. Unsupported functions, unavailable CUDA and other devices still produce IO errors. GPU source recognition happens during Lean elaboration; native compilation and execution remain external trust boundaries. The compilation argument is supplied automatically.

def Function.Internal.record {α : Type} [TorchLean.Storage α] [Runtime.Autograd.Torch.TensorTransfer α] {input output : Spec.Shape} (recording : TorchLean.autograd.Internal.Recording α input output) (value : TorchLean.Tensor α input) (device : NN.Backend.Device) : IO (TorchLean.autograd.Result α input output)

Record with the existing tape without narrowing a custom function's element carrier.

def Function.Internal.tensor {α : Type} [TorchLean.Storage α] [Runtime.Autograd.Torch.TensorTransfer α] {input output : Spec.Shape} (f : TorchLean.Tensor α input → TorchLean.Tensor α output) (recording : Except String (TorchLean.autograd.Internal.Recording α input output)) (value : TorchLean.Tensor α input) (device : NN.Backend.Device) : IO (TorchLean.Tensor α output)

Whole-tensor GPU execution shares the recorded runtime's precision policy.

def NN.Kernel.Frontend.tacticPrepare_function_execution : Lean.ParserDescr

Infer scalar mapping or whole-tensor execution and prepare the existing operation recording.

def Function.run {α ι ο : Type} [TorchLean.Storage α] {input output : Spec.Shape} (f : ι → ο) (value : TorchLean.Tensor α input) [Internal.Signature α ι ο input output] (device : NN.Backend.Device := NN.Backend.Device.cpu) (grad : Bool := false) (execution : Internal.Execution f α input output grad := by prepare_function_execution) : IO (match grad with | false => TorchLean.Tensor α output | true => TorchLean.autograd.Result α input output)

Execute an ordinary scalar or tensor function with inferred shapes and element type.

Scalar functions map over entries; tensor functions receive the whole input. CPU forward-only calls evaluate the original function. Set grad := true to record supported operations on the existing tape and return a value with backward and close. Unsupported recordings fail rather than receive invented derivatives. Native Float32 and Float keep binary32 and binary64 on the same GPU tape. Configured binary arithmetic retains complete words on that tape; its operator subset is narrower than native LibTorch's. CPU-only carriers keep their precision through a reported fallback. Scalar branch decisions synchronize to Lean and record only the chosen path. The execution argument is prepared automatically at the call site.

def Function.zip {α : Type} [TorchLean.Storage α] {shape : Spec.Shape} (f : α → α → α) (left right : TorchLean.Tensor α shape) (device : NN.Backend.Device := NN.Backend.Device.cpu) (gpuScalar? : Option (NN.Kernel.Scalar α) := by first | exact some inferInstance | exact none) (compilation : Except String (NN.Kernel.Zip f) := by prepare_tensor_function) : IO (TorchLean.Tensor α shape)

Combine corresponding entries of two tensors on the selected device.

The shared shape prevents truncation or accidental broadcasting. CPU uses the packed binary map; GPU compiles the binary scalar function into one custom operation. Captured scalar parameters are allowed, with the same source-equivalence and foreign-code boundaries as Function.run.