Adding complex to complex FFT implementation (#4903)

* added cfft function

* fixed formatting issues

* Apply suggestion from Copilot

* Update documentation for cfft

* Remove duplicated logic in stft

* Improve test coverage for cfft

* Fix formatting issues

* Fix import statement

* Update crates/burn-backend-tests/tests/cubecl/fft.rs

* Update crates/burn-tensor/src/tensor/signal/fft.rs

Author: RunjiaChen

crates/burn-backend-tests/tests/cubecl/fft.rs

@@ -1,5 +1,5 @@
 use super::*;
-use burn_tensor::signal::{irfft, rfft};
+use burn_tensor::signal::{cfft, irfft, rfft};
 use burn_tensor::{TensorData, Tolerance};
 
 #[test]
@@ -468,3 +468,243 @@ fn rfft_2d_with_n_padded() {
         );
     }
 }
+
+// ---- cfft tests ----
+
+#[test]
+fn cfft_output_has_n_bins() {
+    // cfft should return N bins, not N/2+1
+    let re = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]);
+    let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]);
+    let (out_re, out_im) = cfft(re, im, 0, None);
+
+    assert_eq!(out_re.dims(), [4]);
+    assert_eq!(out_im.dims(), [4]);
+}
+
+#[test]
+fn cfft_pure_real_input() {
+    // When imaginary part is zero, cfft should produce the same result
+    // as extending rfft to the full spectrum
+    let signal = [1.0f32, 2.0, 3.0, 4.0];
+    let re = TestTensor::<1>::from(signal);
+    let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, None);
+
+    // Expected: DFT of [1,2,3,4]
+    // X[0] = 10, X[1] = -2+2i, X[2] = -2, X[3] = -2-2i
+    let expected_re = TensorData::from([10.0, -2.0, -2.0, -2.0]);
+    let expected_im = TensorData::from([0.0, 2.0, 0.0, -2.0]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_re, Tolerance::absolute(1e-3));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_im, Tolerance::absolute(1e-3));
+}
+
+#[test]
+fn cfft_pure_imaginary_input() {
+    // Signal is purely imaginary: z[n] = i * [1, 2, 3, 4]
+    // FFT(i*x) = i*FFT(x), so result_re = -FFT(x)_im, result_im = FFT(x)_re
+    let re = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]);
+    let im = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, None);
+
+    // FFT([1,2,3,4]) = [10, -2+2i, -2, -2-2i]
+    // i * FFT(x) = i * [10, -2+2i, -2, -2-2i]
+    //            = [-0, -2+(-2)i, 0, 2+(-2)i]  → re = [0, -2, 0, 2], im = [10, -2, -2, -2]
+    let expected_re = TensorData::from([0.0, -2.0, 0.0, 2.0]);
+    let expected_im = TensorData::from([10.0, -2.0, -2.0, -2.0]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_re, Tolerance::absolute(1e-3));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_im, Tolerance::absolute(1e-3));
+}
+
+#[test]
+fn cfft_complex_exponential() {
+    // z[n] = exp(i * 2π * n / 4) for n=0..3, i.e. frequency bin 1
+    // re = [cos(0), cos(π/2), cos(π), cos(3π/2)] = [1, 0, -1, 0]
+    // im = [sin(0), sin(π/2), sin(π), sin(3π/2)] = [0, 1, 0, -1]
+    // DFT should be: X[0]=0, X[1]=4, X[2]=0, X[3]=0
+    let re = TestTensor::<1>::from([1.0, 0.0, -1.0, 0.0]);
+    let im = TestTensor::<1>::from([0.0, 1.0, 0.0, -1.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, None);
+
+    let expected_re = TensorData::from([0.0, 4.0, 0.0, 0.0]);
+    let expected_im = TensorData::from([0.0, 0.0, 0.0, 0.0]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_re, Tolerance::absolute(1e-3));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_im, Tolerance::absolute(1e-3));
+}
+
+#[test]
+fn cfft_zeros() {
+    let re = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]);
+    let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, None);
+
+    let expected = TensorData::from([0.0, 0.0, 0.0, 0.0]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected, Tolerance::absolute(1e-4));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected, Tolerance::absolute(1e-4));
+}
+
+#[test]
+fn cfft_dim1_2d_tensor() {
+    // Apply cfft along dim=1 on a 2D tensor
+    // Row 0: pure real [1, 2, 3, 4] → DFT = [10, -2+2i, -2, -2-2i]
+    // Row 1: complex exponential exp(i·2π·n/4) → DFT = [0, 4, 0, 0]
+    let re = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0], [1.0, 0.0, -1.0, 0.0]]);
+    let im = TestTensor::<2>::from([[0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, -1.0]]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 1, None);
+
+    // Output should be [2, 4] (N=4 bins per row)
+    assert_eq!(cfft_re.dims(), [2, 4]);
+    assert_eq!(cfft_im.dims(), [2, 4]);
+
+    // Row 0: DFT of [1,2,3,4]+i*0 = [10, -2+2i, -2, -2-2i]
+    // Row 1: DFT of exp(i*2π*n/4) = [0, 4, 0, 0]
+    let expected_re = TensorData::from([[10.0, -2.0, -2.0, -2.0], [0.0, 4.0, 0.0, 0.0]]);
+    let expected_im = TensorData::from([[0.0, 2.0, 0.0, -2.0], [0.0, 0.0, 0.0, 0.0]]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_re, Tolerance::absolute(1e-3));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_im, Tolerance::absolute(1e-3));
+}
+
+#[test]
+fn cfft_with_n_padding() {
+    // Signal length 2, padded to N=4
+    // z = [1+0i, 0+0i] padded to [1+0i, 0, 0, 0]
+    // DFT = [1, 1, 1, 1] (all real, zero imag)
+    let re = TestTensor::<1>::from([1.0, 0.0]);
+    let im = TestTensor::<1>::from([0.0, 0.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, Some(4));
+
+    assert_eq!(cfft_re.dims(), [4]);
+
+    let expected_re = TensorData::from([1.0, 1.0, 1.0, 1.0]);
+    let expected_im = TensorData::from([0.0, 0.0, 0.0, 0.0]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_re, Tolerance::absolute(1e-3));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_im, Tolerance::absolute(1e-3));
+}
+
+#[test]
+fn cfft_length_1() {
+    // N=1: DFT of a single complex value is itself
+    let re = TestTensor::<1>::from([3.0]);
+    let im = TestTensor::<1>::from([5.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, None);
+
+    assert_eq!(cfft_re.dims(), [1]);
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&TensorData::from([3.0]), Tolerance::absolute(1e-4));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&TensorData::from([5.0]), Tolerance::absolute(1e-4));
+}
+
+#[test]
+fn cfft_length_2() {
+    // N=2: z = [a, b] → X[0] = a+b, X[1] = a-b
+    // z = [1+2i, 3+4i]
+    // X[0] = (1+3) + i(2+4) = 4+6i
+    // X[1] = (1-3) + i(2-4) = -2-2i
+    let re = TestTensor::<1>::from([1.0, 3.0]);
+    let im = TestTensor::<1>::from([2.0, 4.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, None);
+
+    assert_eq!(cfft_re.dims(), [2]);
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&TensorData::from([4.0, -2.0]), Tolerance::absolute(1e-4));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&TensorData::from([6.0, -2.0]), Tolerance::absolute(1e-4));
+}
+
+#[test]
+#[should_panic(expected = "same shape")]
+fn cfft_rejects_mismatched_shapes() {
+    let re = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]);
+    let im = TestTensor::<1>::from([1.0, 2.0]);
+    let _ = cfft(re, im, 0, None);
+}
+
+#[test]
+fn cfft_dim0_2d_tensor() {
+    // Apply cfft along dim=0 on a 2D tensor (4 rows, 2 columns)
+    // Column 0: complex exponential exp(i·2π·n/4) → DFT = [0, 4, 0, 0]
+    // Column 1: pure real [1, 2, 3, 4] → DFT = [10, -2+2i, -2, -2-2i]
+    let re = TestTensor::<2>::from([[1.0, 1.0], [0.0, 2.0], [-1.0, 3.0], [0.0, 4.0]]);
+    let im = TestTensor::<2>::from([[0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [-1.0, 0.0]]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, None);
+
+    assert_eq!(cfft_re.dims(), [4, 2]);
+    assert_eq!(cfft_im.dims(), [4, 2]);
+
+    let expected_re = TensorData::from([[0.0, 10.0], [4.0, -2.0], [0.0, -2.0], [0.0, -2.0]]);
+    let expected_im = TensorData::from([[0.0, 0.0], [0.0, 2.0], [0.0, 0.0], [0.0, -2.0]]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_re, Tolerance::absolute(1e-3));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_im, Tolerance::absolute(1e-3));
+}
+
+#[test]
+fn cfft_with_n_truncation() {
+    // Signal length 8, truncated to n=4 → DFT of [1+0i, 2+0i, 3+0i, 4+0i]
+    // Trailing values are discarded, not included in the transform.
+    let re = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0, 99.0, 99.0, 99.0, 99.0]);
+    let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0, 99.0, 99.0, 99.0, 99.0]);
+
+    let (cfft_re, cfft_im) = cfft(re, im, 0, Some(4));
+
+    assert_eq!(cfft_re.dims(), [4]);
+
+    // DFT of [1,2,3,4] = [10, -2+2i, -2, -2-2i]
+    let expected_re = TensorData::from([10.0, -2.0, -2.0, -2.0]);
+    let expected_im = TensorData::from([0.0, 2.0, 0.0, -2.0]);
+
+    cfft_re
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_re, Tolerance::absolute(1e-3));
+    cfft_im
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&expected_im, Tolerance::absolute(1e-3));
+}

crates/burn-tensor/src/tensor/signal/fft.rs

@@ -1,3 +1,4 @@
+use alloc::vec;
 use burn_backend::Backend;
 
 use crate::Tensor;
@@ -149,3 +150,118 @@ pub fn irfft<B: Backend, const D: usize>(
     );
     Tensor::new(TensorPrimitive::Float(signal))
 }
+
+/// Computes the 1-dimensional discrete Fourier Transform of complex-valued input.
+///
+/// Internally calls [`rfft`] on the real and imaginary parts separately,
+/// extends each half-spectrum to the full `N`-bin spectrum via Hermitian
+/// symmetry.
+///
+/// Autodiff is not yet supported.
+///
+#[cfg_attr(
+    doc,
+    doc = r#"
+
+Due to the linearity of the Fourier Transform, a complex-valued signal $x\[n\] = x_{re}\[n\] + i x_{im}\[n\]$ can be transformed by applying the FFT to its real and imaginary parts separately:
+
+$$ \text{FFT}(x\[n\]) = \text{FFT}(x_{re}\[n\]) + i \text{FFT}(x_{im}\[n\]) $$
+
+Since $x_{re}\[n\]$ and $x_{im}\[n\]$ are purely real, their transforms can be computed efficiently using the real FFT ([`rfft`]). The full spectrum is then reconstructed by exploiting Hermitian symmetry.
+"#
+)]
+#[cfg_attr(not(doc), doc = r"X\[k\] = Σ x\[n\] * exp(-i*2πkn/N)")]
+///
+/// # Arguments
+///
+/// * `signal_re` - The real part of the complex input signal.
+/// * `signal_im` - The imaginary part of the complex input signal. Must have the
+///   same shape as `signal_re`.
+/// * `dim` - The dimension along which to take the FFT.
+/// * `n` - Optional FFT length. When `None`, the signal must be a power of two
+///   along `dim`. When `Some(n)`, `n` must also be a power of two; the signal is
+///   truncated or zero-padded to length `n`.
+///
+/// # Returns
+///
+/// A tuple `(re, im)` representing the full complex spectrum, each with `n`
+/// elements along `dim`.
+///
+/// # Example
+///
+/// ```rust
+/// use burn_tensor::backend::Backend;
+/// use burn_tensor::Tensor;
+///
+/// fn example<B: Backend>() {
+///     let device = B::Device::default();
+///     let re = Tensor::<B, 1>::from_floats([1.0, 0.0, -1.0, 0.0], &device);
+///     let im = Tensor::<B, 1>::from_floats([0.0, 1.0, 0.0, -1.0], &device);
+///     let (spec_re, spec_im) = burn_tensor::signal::cfft(re, im, 0, None);
+/// }
+/// ```
+pub fn cfft<B: Backend, const D: usize>(
+    signal_re: Tensor<B, D>,
+    signal_im: Tensor<B, D>,
+    dim: usize,
+    n: Option<usize>,
+) -> (Tensor<B, D>, Tensor<B, D>) {
+    assert!(
+        signal_re.shape() == signal_im.shape(),
+        "cfft: signal_re and signal_im must have the same shape, \
+         got {:?} and {:?}",
+        signal_re.shape(),
+        signal_im.shape(),
+    );
+
+    check!(TensorCheck::check_dim::<D>(dim));
+    let fft_size = n.unwrap_or(signal_re.dims()[dim]);
+
+    // rfft validates power-of-two and n constraints internally
+    let (xr, xi) = rfft(signal_re, dim, n);
+    let (yr, yi) = rfft(signal_im, dim, n);
+
+    // Extend half-spectra (N/2+1 bins) to full N-bin spectra via Hermitian symmetry
+    let (xr, xi) = hermitian_extend(xr, xi, dim, fft_size);
+    let (yr, yi) = hermitian_extend(yr, yi, dim, fft_size);
+
+    // FFT(z) = FFT(x) + i·FFT(y)
+    //        = (Xr + i·Xi) + i·(Yr + i·Yi)
+    //        = (Xr - Yi) + i·(Xi + Yr)
+    (xr - yi, xi + yr)
+}
+
+/// Extend a half-spectrum from [`rfft`] (`N/2 + 1` bins) to the full `N`-bin
+/// spectrum using Hermitian symmetry: `X[k] = conj(X[N-k])` for `k > N/2`.
+pub(super) fn hermitian_extend<B: Backend, const D: usize>(
+    half_re: Tensor<B, D>,
+    half_im: Tensor<B, D>,
+    dim: usize,
+    full_len: usize,
+) -> (Tensor<B, D>, Tensor<B, D>) {
+    let half_len = half_re.dims()[dim]; // N/2 + 1
+
+    // For N <= 2, the half-spectrum already covers all bins
+    if full_len <= half_len {
+        return (half_re, half_im);
+    }
+
+    // Mirror bins: reverse of bins 1..N/2-1 (skipping the Nyquist bin),
+    // with conjugated imaginary part. This produces X[N/2+1], X[N/2+2], ..., X[N-1]
+    let mirror_len = full_len - half_len; // N/2 - 1
+    let mirror_re = half_re
+        .clone()
+        .narrow(dim, 1, mirror_len)
+        .flip([dim as isize]);
+    let mirror_im = half_im
+        .clone()
+        .narrow(dim, 1, mirror_len)
+        .flip([dim as isize])
+        .neg();
+
+    // Full spectrum = [half_spectrum, conjugate_mirror]
+    let full_re = Tensor::cat(vec![half_re, mirror_re], dim);
+    let full_im = Tensor::cat(vec![half_im, mirror_im], dim);
+
+    (full_re, full_im)
+}