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.

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 (FloatLib.Floats.ExecFloat.Binary 8 23), selected via --arithmetic=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 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.fftLeadingAxis (n : ℕ) (rest : Spec.Shape) : Layer (rest.prependDim n) (rest.prependDim n)

Implementation of 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.

The public fftAtDepth operation moves an arbitrary axis here and restores the original axis order.

def Runtime.Autograd.Model.Layers.FFT.Internal.permuteBySwaps {α : Type} [TorchLean.Storage α] [Context α] {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 the outer-axis implementation, then swapping back.

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

def Runtime.Autograd.Model.Layers.FFT.fftAtDepth {s : Spec.Shape} : ℕ → Layer s s