FFT (1D) building blocks

TorchLean’s layer/model definitions are scalar-polymorphic: a model runs over whatever scalar type α you instantiate it with (e.g. Float, IEEE32Exec, ℝ, etc.). A “real FFT” would normally change the scalar type (real → complex), but TorchLean’s LayerDef does not support changing the scalar type mid-model.

So this module provides complex-domain transforms: fft and ifft as layers that assume the chosen scalar type α already behaves like a complex field (for example TorchLean.Complex IEEE32Exec, selected via --dtype=complex).

Implementation note: we define fft/ifft as multiplication by explicit DFT matrices (so they are purely built from existing ops like const and matmul). This is correctness-first and keeps the transform differentiable under the existing autograd rules. It is not optimized for large n.

Numerics note:

• Over mathlib’s ℂ, the corresponding DFT/IDFT inversion facts are proved in NN.Proofs.Analysis.Fft (and the bridge to these twiddle/matrix definitions is in NN.Proofs.Analysis.FftBridge).
• For executable IEEE32Exec, sin/cos are implemented deterministically in Lean (see NN.Floats.IEEEExec.Exec32). This makes FFT execution reproducible across platforms. Proving tight end-to-end accuracy bounds for FFT still requires a separate analysis layer (or an interval/oracle backend) to relate those executable trigonometric approximations to real sin/cos.

@[reducible, inline]

abbrev Runtime.Autograd.TorchLean.NN.FFT1D.vec (n : ℕ) : Spec.Shape

Shape abbreviation for a length-n vector.

@[reducible, inline]

abbrev Runtime.Autograd.TorchLean.NN.FFT1D.mat (m n : ℕ) : Spec.Shape

Shape abbreviation for an m×n matrix.

We build twiddle factors using only the Context interface: I := sqrt(-1) and e^{-iθ} = cos θ - I * sin θ.

This is intended to be instantiated with TorchLean’s own complex scalar TorchLean.Complex β (for some base scalar β). For real-only scalar backends, the formulas are not meaningful.

def Runtime.Autograd.TorchLean.NN.FFT1D.I {α : Type} [Context α] : α

The imaginary unit, represented as sqrt(-1) in the ambient scalar type.

def Runtime.Autograd.TorchLean.NN.FFT1D.twiddle {α : Type} [Context α] (n j k : ℕ) : α

Twiddle factor exp(-2π i * j*k / n) written as cos θ - i sin θ.

def Runtime.Autograd.TorchLean.NN.FFT1D.twiddleInv {α : Type} [Context α] (n j k : ℕ) : α

Twiddle factor exp(+2π i * j*k / n) written as cos θ + i sin θ.

def Runtime.Autograd.TorchLean.NN.FFT1D.dftMatrix {α : Type} [Context α] (n : ℕ) : Spec.Tensor α (mat n n)

DFT matrix F : n×n with entries F[k,j] = exp(-2π i j*k / n).

def Runtime.Autograd.TorchLean.NN.FFT1D.idftMatrix {α : Type} [Context α] (n : ℕ) : Spec.Tensor α (mat n n)

Inverse DFT matrix F⁻¹ : n×n with entries F⁻¹[j,k] = exp(+2π i j*k / n) / n.

def Runtime.Autograd.TorchLean.NN.FFT1D.fftDim0 (n : ℕ) (rest : Spec.Shape) : LayerDef (Spec.Shape.dim n rest) (Spec.Shape.dim n rest)

FFT along the outermost axis of a tensor.

This applies the DFT to the leading dimension n of a shape dim n rest by:

1. reshaping to a matrix n × (numel rest),
2. left-multiplying by the n×n DFT matrix, then
3. reshaping back.

This is the most generally useful primitive for building N-D FFTs: you can permute an axis to the front, call fftDim0, then permute back.

def Runtime.Autograd.TorchLean.NN.FFT1D.ifftDim0 (n : ℕ) (rest : Spec.Shape) : LayerDef (Spec.Shape.dim n rest) (Spec.Shape.dim n rest)

Inverse FFT along the outermost axis of a tensor (uses the inverse DFT matrix).

See fftDim0 for the implementation strategy.

@[reducible, inline]

abbrev Runtime.Autograd.TorchLean.NN.FFT1D.fftMat (n width : ℕ) : LayerDef (mat n width) (mat n width)

FFT on matrices: apply the DFT along the leading dimension (n×width).

@[reducible, inline]

abbrev Runtime.Autograd.TorchLean.NN.FFT1D.ifftMat (n width : ℕ) : LayerDef (mat n width) (mat n width)

Inverse FFT on matrices: apply the inverse DFT along the leading dimension (n×width).

@[reducible, inline]

abbrev Runtime.Autograd.TorchLean.NN.FFT1D.fftVec (n : ℕ) : LayerDef (vec n) (vec n)

FFT on vectors.

@[reducible, inline]

abbrev Runtime.Autograd.TorchLean.NN.FFT1D.ifftVec (n : ℕ) : LayerDef (vec n) (vec n)

Inverse FFT on vectors.

def Runtime.Autograd.TorchLean.NN.FFT1D.Internal.permuteBySwaps {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (x : (s : Spec.Shape) × RefTy m α s) (swaps : List ℕ) : m ((s' : Spec.Shape) × RefTy m α s')

Apply a sequence of swapAdjacentAtDepth operations (shape-indexed permutation primitive).

FFT along an axis at a given depth (0-based from the outermost).

This is implemented by swapping the target axis outward (one adjacent swap per step) until it reaches depth 0, applying fftDim0, then swapping back.

If depth ≥ Shape.rank s, this layer is the identity.

def Runtime.Autograd.TorchLean.NN.FFT1D.fftAtDepth {s : Spec.Shape} : ℕ → LayerDef s s

def Runtime.Autograd.TorchLean.NN.FFT1D.ifftAtDepth {s : Spec.Shape} : ℕ → LayerDef s s

Inverse FFT along an axis at a given depth (see fftAtDepth).