Fourier transforms along a selected axis

TorchLean’s layer/model definitions are scalar-polymorphic: a model runs over whatever scalar type $\alpha$ you instantiate it with (for example Float, ExecFloat.Binary 8 23, or $\mathbb{R}$). A “real FFT” would normally change the scalar type (real $\to$ complex), but TorchLean’s Layer does not support changing the scalar type mid-model.

This module provides DFT and inverse DFT matrices, and a forward dft layer that assumes $\alpha$ already behaves like a complex field (for example TorchLean.Complex (FloatLib.Floats.ExecFloat.Binary 8 23)).

The layer multiplies by an explicit DFT matrix using const and matmul, so the existing autograd rules apply. This is a dense transform, not a fast FFT algorithm; it is not intended for large n. The inverse matrix does not currently have a corresponding layer constructor here.

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 ExecFloat.Binary 8 23, sin/cos are implemented deterministically in Lean (see FloatLib.Floats.Formats.BinaryInterchange.Configured.Transcendentals). 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.

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.Model.Layers.FFT.I {α : Type} [TorchLean.Storage α] [Context α] : α

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

def Runtime.Autograd.Model.Layers.FFT.twiddle {α : Type} [TorchLean.Storage α] [Context α] (n j k : ℕ) : α

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

def Runtime.Autograd.Model.Layers.FFT.twiddleInv {α : Type} [TorchLean.Storage α] [Context α] (n j k : ℕ) : α

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

def Runtime.Autograd.Model.Layers.FFT.dftMatrix {α : Type} [TorchLean.Storage α] [Context α] (n : ℕ) : TorchLean.Tensor α [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.Model.Layers.FFT.idftMatrix {α : Type} [TorchLean.Storage α] [Context α] (n : ℕ) : TorchLean.Tensor α [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.Model.Layers.FFT.Internal.dft (n : ℕ) (rest : Spec.Shape) : Layer (rest.prependDim n) (rest.prependDim n)

Dense Fourier transform 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.

The public dft operation moves the selected axis here and restores the original axis order.

def Runtime.Autograd.Model.Layers.FFT.dft {s : Spec.Shape} (axis : ℕ) : Layer s s

Dense Fourier transform along axis, numbered from zero.

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

If $\mathtt{axis}\ge\operatorname{rank}(s)$, this layer is the identity.