feat(libtorch): piecewise-linear table lookup composed from float32 primitives

NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean

@@ -477,6 +477,10 @@ opaque abs (b : @& Buffer) : Buffer
 @[never_extract, extern "torchlean_cuda_buffer_abs_bwd"]
 opaque absBwd (x dLdy : @& Buffer) : Buffer
 
+/-- Elementwise floor: the largest integral value not above each element, as `Float32.floor`. -/
+@[never_extract, extern "torchlean_cuda_buffer_floor"]
+opaque floor (b : @& Buffer) : Buffer
+
 /-- Elementwise `sqrt (max x 0)`, matching `Tensor.sqrtSpec`.
 
 Negative inputs and either signed zero return positive zero. NaN inputs remain NaN; the clamp

NN/Runtime/Autograd/Engine/LibTorch/KernelSpec.lean

@@ -298,6 +298,23 @@ def scatterAddSpec {n k : Nat} (x : FlatBuffer n) (values : FlatBuffer k)
       (fun acc j => if idx j = i then ExecFloat.add acc (values j) else acc)
       (x i)
 
+/-! ## Table lookup -/
+
+/--
+Piecewise-linear table lookup reference: `v₀ + (v₁ - v₀) · w` with one rounding per operation.
+
+`lower` and `upper` are the sample positions each element interpolates between and `weight` is
+its interpolation weight. `Buffer.tableLookupForward` composes exactly this sequence of float32
+primitives, so a native result that agrees with this spec bit for bit inherits the scalar facts
+of `Float32Contract`.
+-/
+def tableLookupSpec {size k : Nat} (table : FlatBuffer size)
+    (lower upper : Fin k → Fin size) (weight : FlatBuffer k) : FlatBuffer k :=
+  fun i =>
+    let lowerSample := table (lower i)
+    ExecFloat.add lowerSample
+      (ExecFloat.mul (ExecFloat.sub (table (upper i)) lowerSample) (weight i))
+
 /-! ## Batched row-major matrix multiplication -/
 
 /-- Linear index for either BMM input, viewed as row-major `(batch, rows, cols)` storage. -/

NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean

@@ -363,6 +363,29 @@ Indices that fit in `UInt32` but are out of bounds are ignored; repeated indices
 opaque scatterAddRows (mat values : @& Buffer) (rows cols : UInt32) (indices : @& Array Nat)
   (k : UInt32) : Buffer
 
+/--
+Gather `table[positions[i]]` for `count` device-resident positions.
+
+`table` has `size > 0` entries. Positions are float32 values computed on device: a NaN reads as
+`0`, every position is clamped to `[0, size - 1]` and then truncated, so a non-integral position
+floors and no selection leaves the table. The clamp is applied to the float32 position and again
+to the integer index, so it holds for tables above 2^24 entries, where `size - 1` has no float32
+representation; which entries a float32 position can address exactly is a separate question,
+bounded by `Tape.tableLookup`. The result has `count` entries.
+-/
+@[never_extract, extern "torchlean_cuda_buffer_gather_at"]
+opaque gatherAt (table : @& Buffer) (size : UInt32) (positions : @& Buffer) (count : UInt32) :
+  Buffer
+
+/--
+Scatter-add `count` values into a fresh zero buffer of `size > 0` entries at device-resident
+positions, addressed exactly as in `gatherAt`: the VJP of `gatherAt` with respect to its table.
+Repeated positions accumulate; the native accumulation order is unspecified unless deterministic
+algorithms are selected.
+-/
+@[never_extract, extern "torchlean_cuda_buffer_scatter_add_at"]
+opaque scatterAddAt (size : UInt32) (positions values : @& Buffer) (count : UInt32) : Buffer
+
 end Buffer
 
 end LibTorch

NN/Runtime/Autograd/Engine/LibTorch/Ops.lean

@@ -15,6 +15,7 @@ public import NN.Runtime.Autograd.Engine.LibTorch.Ops.Linear
 public import NN.Runtime.Autograd.Engine.LibTorch.Ops.NormSoftmax
 public import NN.Runtime.Autograd.Engine.LibTorch.Ops.SelectiveScan
 public import NN.Runtime.Autograd.Engine.LibTorch.Ops.Shape
+public import NN.Runtime.Autograd.Engine.LibTorch.Ops.TableLookup
 
 /-!
 CUDA operation dispatch surface.

NN/Runtime/Autograd/Engine/LibTorch/Ops/Elementwise.lean

@@ -68,6 +68,12 @@ returned buffers are owned by the tape/gradient accumulator; workspace buffers a
     (forward := Buffer.abs)
     (backward := fun x dLdy => Buffer.absBwd x dLdy)
 
+/-- Pointwise floor node. The function is piecewise constant, so its backward is zero. -/
+@[inline] def floor {s : Shape} (t : Tape) (xId : Nat) : Result (Tape × Nat) :=
+  unary (t := t) "floor" xId s s
+    (forward := Buffer.floor)
+    (backward := fun _x dLdy => Buffer.zeros (Buffer.size dLdy))
+
 /-- Pointwise square-root node using the CUDA buffer derivative convention. -/
 @[inline] def sqrt {s : Shape} (t : Tape) (xId : Nat) : Result (Tape × Nat) :=
   unary (t := t) "sqrt" xId s s

NN/Runtime/Autograd/Engine/LibTorch/Ops/Indexing.lean

@@ -226,6 +226,38 @@ end Indexing
           (sourceId, { s := sourceShape, buf := dSource })] }
   pure (t.addNode node)
 
+/-!
+## Device-resident positions
+-/
+
+/--
+Gather entries of a rank-one table at positions held in another tape value.
+
+The positions are float32 values computed on device and addressed as `Buffer.gatherAt` documents.
+The result takes the positions' shape. The table receives the scatter-add VJP; the positions are
+selectors rather than differentiable inputs, so no gradient flows to them.
+-/
+@[inline] def gatherAt {size : Nat} {s : Shape} (t : Tape) (tableId positionsId : Nat) :
+    Result (Tape × Nat) := do
+  if size = 0 then
+    throw "autograd: gather_at: the table is empty"
+  let tableShape : Shape := .dim size .scalar
+  let table ← requireValue (t := t) tableId tableShape
+  let positions ← requireValue (t := t) positionsId s
+  let size32 ← AnyBuffer.natToU32Checked size
+  let count32 ← AnyBuffer.numelU32 s
+  let y := Buffer.gatherAt table size32 positions count32
+  let node : Node :=
+    { name := some "gather_at"
+      value := { s := s, buf := y }
+      requiresGrad := (t.getNode? tableId).any (·.requiresGrad)
+      parents := #[tableId, positionsId]
+      backward := fun dLdyAny => do
+        let dLdy ← requireGrad dLdyAny s
+        let dTable := Buffer.scatterAddAt size32 positions dLdy.buf count32
+        pure #[(tableId, { s := tableShape, buf := dTable })] }
+  pure (t.addNode node)
+
 end Tape
 
 end LibTorch

NN/Runtime/Autograd/Engine/LibTorch/Ops/TableLookup.lean

@@ -0,0 +1,178 @@
+/-
+Copyright (c) 2026 TorchLean
+Released under MIT license as described in the file LICENSE.
+Authors: TorchLean Team
+-/
+
+module
+
+public import NN.Runtime.Autograd.Engine.LibTorch.Ops.Core
+
+/-!
+# CUDA Tape Operations: Piecewise-Linear Table Lookup
+
+A function tabulated at `width` uniformly spaced grid points is evaluated at a grid coordinate
+`u` by linear interpolation between the two samples around `u`, after clamping `u` to
+`[0, width - 1]`. This is the uniform-grid analogue of `np.interp`, which is how calibration
+curves, inverted sensor-response tables, and activation lookups are evaluated. A layered table
+stores one such function per layer and selects each element's layer from a second buffer.
+
+With `i = ⌊u⌋` and `w = u - i`, the value is `v₀ + (v₁ - v₀) · w` for the samples `v₀ = table[i]`
+and `v₁ = table[min (i + 1) (width - 1)]`. The forward composes numerical primitives that each
+round once in float32, so the result is `KernelSpec.tableLookupSpec` bit for bit. The backward is
+TorchLean's own: the coordinate receives the slope of its segment, zero where the clamp is active
+as for `clamp`, and the table receives the two interpolation weights scattered to the samples
+they multiplied.
+-/
+
+@[expose] public section
+
+namespace Runtime
+namespace Autograd
+namespace LibTorch
+
+open Spec TorchLean
+
+namespace Buffer
+
+/-- The samples each element interpolates between, computed once and reused by the backward. -/
+structure TableSegments where
+  /-- Flat position of the lower sample. -/
+  lower : Buffer
+  /-- Flat position of the upper sample; the lower one again on the last node of a layer. -/
+  upper : Buffer
+  /-- Interpolation weight toward the upper sample, in `[0, 1)`. -/
+  weight : Buffer
+  /-- Intermediates that the forward value releases once it has consumed them. -/
+  workspace : Array Buffer := #[]
+
+/--
+Locate the segment of each of `count` coordinates inside a table of `layers × width` samples.
+
+Coordinates clamp to `[0, width - 1]` and the optional layer indices to `[0, layers - 1]` with the
+backend's `clamp`, whose NaN selects the lower bound. Every operation here is exact in float32
+while `layers * width ≤ 2^24`; the tape operation enforces that bound.
+-/
+@[no_expose] def tableSegments (layers width : UInt32) (coords : Buffer) (layer : Option Buffer)
+    (count : UInt32) : TableSegments :=
+  let last : Float := Float.ofNat (width.toNat - 1)
+  let coordinate := clamp coords 0.0 last
+  let column := floor coordinate
+  let weight := releaseThen coordinate (sub coordinate column)
+  let ones := full count 1.0
+  let next := releaseThen ones (add column ones)
+  let nextColumn := releaseThen next (clamp next 0.0 last)
+  match layer with
+  | none => { lower := column, upper := nextColumn, weight := weight }
+  | some layerIdx =>
+      let lastLayer : Float := Float.ofNat (layers.toNat - 1)
+      let clampedLayer := clamp layerIdx 0.0 lastLayer
+      let integralLayer := releaseThen clampedLayer (floor clampedLayer)
+      let offset := releaseThen integralLayer (scale integralLayer (Float.ofNat width.toNat))
+      let lower := add offset column
+      let upper := add offset nextColumn
+      { lower := lower, upper := upper, weight := weight
+        workspace := #[column, nextColumn, offset] }
+
+/--
+Interpolate the table at located segments: `v₀ + (v₁ - v₀) · w`, one rounding per operation.
+
+The segments' workspace is released through the result, so the forward value is the last use of
+those intermediates.
+-/
+@[no_expose] def tableLookupForward (table : Buffer) (size : UInt32) (seg : TableSegments)
+    (count : UInt32) : Buffer :=
+  let lowerSample := gatherAt table size seg.lower count
+  let upperSample := gatherAt table size seg.upper count
+  let slope := sub upperSample lowerSample
+  let scaled := mul slope seg.weight
+  let value := add lowerSample scaled
+  releaseManyThen (#[lowerSample, upperSample, slope, scaled] ++ seg.workspace) value
+
+/--
+VJP of `tableLookupForward`, returning the coordinate and table cotangents.
+
+The coordinate cotangent is the segment slope times the upstream cotangent, zeroed where the
+coordinate clamp was active (the open-interval rule of `clampBwd`). The table cotangent scatters
+`(1 - w)` and `w` times the upstream cotangent to the lower and upper samples.
+-/
+@[no_expose] def tableLookupBackward (table : Buffer) (size : UInt32) (seg : TableSegments)
+    (coords : Buffer) (width : UInt32) (dLdy : Buffer) (count : UInt32) : Buffer × Buffer :=
+  let last : Float := Float.ofNat (width.toNat - 1)
+  let lowerSample := gatherAt table size seg.lower count
+  let upperSample := gatherAt table size seg.upper count
+  let slope := sub upperSample lowerSample
+  let interior := mul dLdy slope
+  let dCoords := releaseManyThen #[lowerSample, upperSample, slope, interior]
+    (clampBwd coords interior 0.0 last)
+  let ones := full count 1.0
+  let lowerWeight := sub ones seg.weight
+  let dLower := mul dLdy lowerWeight
+  let dUpper := mul dLdy seg.weight
+  let lowerPart := scatterAddAt size seg.lower dLower count
+  let upperPart := scatterAddAt size seg.upper dUpper count
+  let dTable := releaseManyThen #[ones, lowerWeight, dLower, dUpper, lowerPart, upperPart]
+    (add lowerPart upperPart)
+  (dCoords, dTable)
+
+end Buffer
+
+namespace Tape
+
+/-- Shared implementation of both lookups; `layerId` supplies per-element layer indices. -/
+@[no_expose] def tableLookupCore (layers width : Nat) {s : Shape} (t : Tape)
+    (tableId coordsId : Nat) (layerId : Option Nat) : Result (Tape × Nat) := do
+  if width = 0 then
+    throw "autograd: table_lookup: the table has no samples"
+  if layers = 0 then
+    throw "autograd: table_lookup: the table has no layers"
+  if layers * width > 2 ^ 24 then
+    throw "autograd: table_lookup: float32 positions address at most 2^24 samples exactly"
+  let tableShape : Shape := match layerId with
+    | none => .dim width .scalar
+    | some _ => .dim layers (.dim width .scalar)
+  let table ← requireValue (t := t) tableId tableShape
+  let coords ← requireValue (t := t) coordsId s
+  let layer ← match layerId with
+    | none => pure none
+    | some id => some <$> requireValue (t := t) id s
+  let layers32 ← AnyBuffer.natToU32Checked layers
+  let width32 ← AnyBuffer.natToU32Checked width
+  let size32 ← AnyBuffer.natToU32Checked (layers * width)
+  let count32 ← AnyBuffer.numelU32 s
+  let seg := Buffer.tableSegments layers32 width32 coords layer count32
+  let y := Buffer.tableLookupForward table size32 seg count32
+  let node : Node :=
+    { name := some (if layerId.isSome then "layered_table_lookup" else "table_lookup")
+      value := { s := s, buf := y }
+      requiresGrad := (t.getNode? tableId).any (·.requiresGrad) ||
+        (t.getNode? coordsId).any (·.requiresGrad)
+      parents := #[tableId, coordsId] ++ (match layerId with | none => #[] | some id => #[id])
+      cleanup := #[seg.lower, seg.upper, seg.weight]
+      backward := fun dLdyAny => do
+        let dLdy ← requireGrad dLdyAny s
+        let (dCoords, dTable) :=
+          Buffer.tableLookupBackward table size32 seg coords width32 dLdy.buf count32
+        pure #[
+          (tableId, { s := tableShape, buf := dTable }),
+          (coordsId, { s := s, buf := dCoords })] }
+  pure (t.addNode node)
+
+/-- Piecewise-linear lookup of a rank-one table of `width` samples at coordinates of shape `s`. -/
+@[inline] def tableLookup {width : Nat} {s : Shape} (t : Tape) (tableId coordsId : Nat) :
+    Result (Tape × Nat) :=
+  tableLookupCore 1 width (s := s) t tableId coordsId none
+
+/--
+Piecewise-linear lookup of a `layers × width` table, selecting each element's layer from the
+integral values of the tape value `layerId`, which has shape `s` and receives no gradient.
+-/
+@[inline] def layeredTableLookup {layers width : Nat} {s : Shape} (t : Tape)
+    (tableId coordsId layerId : Nat) : Result (Tape × Nat) :=
+  tableLookupCore layers width (s := s) t tableId coordsId (some layerId)
+
+end Tape
+
+end LibTorch
+end Autograd
+end Runtime

NN/Tests/Runtime/Cuda/Suite.lean

@@ -10,6 +10,7 @@ public import NN.Tests.Runtime.ParameterAliases
 public import NN.Tests.Runtime.Cuda.Softmax
 public import NN.Tests.Runtime.Cuda.Elementwise
 public import NN.Tests.Runtime.Cuda.Attention
+public import NN.Tests.Runtime.Cuda.TableLookup
 public import NN.Tests.Runtime.Cuda.SelectiveScan
 public import NN.Tests.Runtime.Cuda.Fft
 public import NN.Tests.Runtime.Cuda.Stress
@@ -37,6 +38,7 @@ def run : IO Unit := do
   Softmax.run
   Elementwise.run
   Attention.run
+  TableLookup.run
   SelectiveScan.run
   Fft.run
   Stress.run