CUDA FNO1D (real RFFT fused path)

This file provides a direct CUDA forward and VJP runner for a real-valued FNO1D model. Its spectral convolution uses the same native tape primitive as nn.models.fnoRfft on CUDA. The public constructor also offers a dense reference with identical one-sided weights.

This runner keeps its explicit buffer lifetime and Adam handling for the Burgers example. Comparisons with the public constructor must load identical parameter tensors and use the same ReLU blocks, loss, and update rule. The arbitrary-rank full-DFT model in Model.Fno has a different spectral parameterization and is not an interchangeable checkpoint.

structure Runtime.Autograd.Cuda.Fno1dRfftFused.Param : Type

Trainable parameter plus Adam moment buffers.

All three tensors share shape. Values are serialized only when uploaded to CUDA; host Adam updates consume downloaded gradients as tensors.

• shape : Spec.Shape

  Runtime tensor shape for value, m, and v.

• value : TorchLean.Tensor Float self.shape

  Current parameter values in row-major order.

• m : TorchLean.Tensor Float self.shape

  Adam first-moment accumulator.

• v : TorchLean.Tensor Float self.shape

  Adam second-moment accumulator.

structure Runtime.Autograd.Cuda.Fno1dRfftFused.Forward : Type

Output of one fused-FNO tape construction.

• tape : Tape

  The completed CUDA tape.

• predId : ℕ

  Node id of the prediction tensor.

• lossId? : Option ℕ

  Optional scalar loss node id, present only when a target was supplied.

• paramIds : Array ℕ

  Tape node ids for parameters, in the same order as the parameter array.

def Runtime.Autograd.Cuda.Fno1dRfftFused.Forward.ownedBufferCount (fw : Forward) : ℕ

Number of CUDA buffer handles owned by a completed forward tape.

def Runtime.Autograd.Cuda.Fno1dRfftFused.Forward.dispose (fw : Forward) : IO Unit

Release every forward value and saved workspace owned by a completed tape.

The fused FNO wrapper rebuilds its tape for every sample. Explicit disposal is therefore part of the wrapper's ownership contract; waiting for external-object finalizers makes long training and evaluation runs retain one full tape per sample until a later runtime collection.

structure Runtime.Autograd.Cuda.Fno1dRfftFused.AdamState : Type

Minimal Adam state carried across fused-FNO training steps.

• step : ℕ

  Step counter (1-based in the Adam bias correction formulas).

• beta1Pow : Float

  Cached beta1^step for bias correction (starts at 1).

• beta2Pow : Float

  Cached beta2^step for bias correction (starts at 1).

def Runtime.Autograd.Cuda.Fno1dRfftFused.uniformAt (seed idx : ℕ) (lo hi : Float) : Float

Deterministic uniform sample in [lo, hi) for a scalar index.

def Runtime.Autograd.Cuda.Fno1dRfftFused.initParam (shape : Spec.Shape) (seed : ℕ) (lo hi : Float) : Param

Initialize a trainable parameter and zero Adam moments.

def Runtime.Autograd.Cuda.Fno1dRfftFused.initBias (shape : Spec.Shape) : Param

Initialize a bias-like parameter at zero with zero Adam moments.

def Runtime.Autograd.Cuda.Fno1dRfftFused.initParams (width modes blocks seed : ℕ) : Array Param

Initialize parameters for the fused FNO1D model:

• input lift: W_in : (1,width), b_in : (width)
• blocks: (wRe,wIm) : (modes,width,width), wSkip : (width,width), bSkip : (width)
• output proj: W_out : (width,1), b_out : (1)

def Runtime.Autograd.Cuda.Fno1dRfftFused.getParam (ps : Array Param) (i : ℕ) : Result Param

Fetch a parameter with an error message that points to the fused-FNO wrapper.

def Runtime.Autograd.Cuda.Fno1dRfftFused.addParamLeaf (t : Tape) (ps : Array Param) (paramBuffers : Array Buffer) (paramIds : Array ℕ) (i : ℕ) : Result (Tape × Array ℕ × ℕ)

Validate uploaded parameter buffer i against its declared shape, add it as a gradient-requiring CUDA tape leaf, and record the new node id.

def Runtime.Autograd.Cuda.Fno1dRfftFused.broadcastVecToMat (t : Tape) (grid cols xId : ℕ) : Result (Tape × ℕ)

Broadcast a vector of length cols across grid rows.

def Runtime.Autograd.Cuda.Fno1dRfftFused.forwardWithBuffers (grid width modes blocks : ℕ) (ps : Array Param) (target? : Option (TorchLean.Tensor Float [grid])) (xBuffer : Buffer) (paramBuffers : Array Buffer) (targetBuffer? : Option Buffer) : Result Forward

Build a CUDA tape that computes prediction (and optionally MSE loss) for the fused real-RFFT FNO.

Inputs:

• x : (grid) (interpreted as (grid,1)),
• optional target : (grid).

Every external CUDA buffer is checked against the corresponding logical shape and element count before an operation can consume it. Input and parameter buffers are checked before their leaf nodes are added; the optional target is checked by the MSE operation on the loss path.

def Runtime.Autograd.Cuda.Fno1dRfftFused.forward (grid width modes blocks : ℕ) (ps : Array Param) (x : TorchLean.Tensor Float [grid]) (target? : Option (TorchLean.Tensor Float [grid])) : IO (Result Forward)

Build one fused-FNO tape from fresh CUDA uploads.

The uploads are effectful so repeated forwards over identical host arrays cannot share an external buffer handle. This gives each returned Forward exclusive ownership of the buffers it disposes.

def Runtime.Autograd.Cuda.Fno1dRfftFused.scalarFromTape (t : Tape) (id : ℕ) : IO (Result Float)

Download a scalar CUDA tape value to host Float.

def Runtime.Autograd.Cuda.Fno1dRfftFused.predFromTape (grid : ℕ) (t : Tape) (id : ℕ) : IO (Result (TorchLean.Tensor Float [grid]))

Download a (grid,1) prediction matrix as a length-grid tensor.

def Runtime.Autograd.Cuda.Fno1dRfftFused.meanLoss (grid width modes blocks : ℕ) (ps : Array Param) (samples : TorchLean.Data.SampleStream (TorchLean.Tensor Float [grid] × TorchLean.Tensor Float [grid])) : IO (Result Float)

Mean MSE over a nonempty indexed sample stream, releasing each tape after scalar download.

def Runtime.Autograd.Cuda.Fno1dRfftFused.adamUpdateBiasCorrected {shape : Spec.Shape} (value m v grad : TorchLean.Tensor Float shape) (lr beta1 beta2 eps biasCorr1 biasCorr2 : Float) : TorchLean.Tensor Float shape × TorchLean.Tensor Float shape × TorchLean.Tensor Float shape

Tensor Adam update with cached bias corrections and the fused path's scalar evaluation order.

Each multiplication/division retains its original ordering, including (1-beta2)*g*g and lr*(m/biasCorr1)/(sqrt(v/biasCorr2)+eps). Equal shapes replace the former runtime length checks.

def Runtime.Autograd.Cuda.Fno1dRfftFused.prepareAdamBackward (fw : Forward) (st : AdamState) (beta1 beta2 : Float) : IO (Result (Tape.SparseGradMap × AdamState × Float × Float))

Run reverse-mode on the fused-FNO tape and update every recorded parameter with Adam.

Gradients are computed on CUDA buffers and downloaded to host arrays before the update. A high-throughput optimizer kernel should live in a separate CUDA optimizer layer, not inside this model helper.

def Runtime.Autograd.Cuda.Fno1dRfftFused.updateParamsAdam (ps : Array Param) (fw : Forward) (lr : Float) (st : AdamState) (beta1 : Float := 0.9) (beta2 : Float := 0.999) (eps : Float := 1e-8) : IO (Result (Array Param × AdamState))

Run one Adam update and deterministically release the consumed gradient and forward buffers.

def Runtime.Autograd.Cuda.Fno1dRfftFused.predict (grid width modes blocks : ℕ) (parameters : Array Param) (input : TorchLean.Tensor Float [grid]) : IO (Result (TorchLean.Tensor Float [grid]))

Predict through the fused spectral path and release the temporary tape after downloading.

def Runtime.Autograd.Cuda.Fno1dRfftFused.trainStep (grid width modes blocks : ℕ) (parameters : Array Param) (input target : TorchLean.Tensor Float [grid]) (learningRate : Float) (state : AdamState) : IO (Result (Array Param × AdamState))

Consume one fused training tape, returning updated parameters and cached Adam state.