FFT (1D) building blocks

TorchLean’s layer/model definitions are scalar-polymorphic: a model runs over whatever scalar type $\alpha$ you instantiate it with (for example Float, IEEE32Exec, or $\mathbb{R}$). A “real FFT” would normally change the scalar type (real $\to$ 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 $\alpha$ 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\theta}=\cos\theta-I\sin\theta$.

This is intended to be instantiated with TorchLean’s own complex scalar TorchLean.Complex β (for some base scalar $\beta$). 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 $e^{-2\pi i jk/n}$ written as $\cos\theta-i\sin\theta$.

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

Twiddle factor $e^{2\pi i jk/n}$ written as $\cos\theta+i\sin\theta$.

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

DFT matrix $F\in\mathbb{C}^{n\times n}$ with entries $F_{k,j}=e^{-2\pi i jk/n}$.

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

Inverse DFT matrix $F^{-1}\in\mathbb{C}^{n\times n}$ with entries $(F^{-1})_{j,k}=e^{2\pi i jk/n}/n$.

def Runtime.Autograd.TorchLean.NN.FFT1D.fftLeadingAxis (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 fftLeadingAxis, then permute back.

def Runtime.Autograd.TorchLean.NN.FFT1D.ifftLeadingAxis (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 fftLeadingAxis 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)

Vector FFT layer, implemented as a DFT along the only non-scalar axis.

@[reducible, inline]

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

Inverse vector FFT layer, implemented as an inverse DFT along the only non-scalar axis.

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 fftLeadingAxis, then swapping back.

If $\mathtt{depth}\ge\operatorname{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).