Repository: google/jax
Issue: linear_transpose involving fft seemingly incorrect

I was planning to implement functionality that requires me to be able to take the transpose of a complicated linear function, which includes fourier transforms. I noticed that the fft module seems to implement the rules for transposition; the last few lines of https://jax.readthedocs.io/en/latest/_modules/jax/_src/lax/fft.html seem to pertain to it.

I am not familiar enough with the jax API at present to point out any bugs; but below is a simple test that demonstrates what I believe to be a bug. The build_matrix helper function explicitly constructs the coefficients of the linear function by feeding it all delta functions in sequence. If I linear_transpose my function, that should be identical to the matrix transpose of the built matrix. Yet it doesnt. It seems as if im getting my result in reverse order (plus another off-by-one index bug I think).

Seems to me like a bug in the implementation of the transpose rules for ffts; but again not qualified myself to spot it.

While on the topic, slightly related question: when viewed as linear operators, convolution and correlation are transposed operators. Should I trust jax to figure out efficient transformations along these lines (assuming the underlying fft rules are bug free); or is it likely optimal for me to figure out how to override the linear transpose of a convolution with my own handcrafted correlation functions (and vice versa)?


Code to reproduce:

```python
import numpy as np

import jax
from jax import numpy as jnp
import matplotlib.pyplot as plt


np.random.seed(0)

signal = np.cumsum(np.random.randn(2**8))
signal_jax = jnp.array(signal)

x = np.linspace(-1, 1, len(signal))
psf = np.clip(0.2 - np.abs(x), 0, 1) * (x > 0)
psf /= psf.sum()
psf_jax = jnp.array(psf)

jrfft = jax.jit(jnp.fft.rfft)
jirfft = jax.jit(jnp.fft.irfft)


@jax.jit
def convolve(a, b):
	fa = jrfft(a)
	fb = jrfft(b)
	return jirfft(fa * fb)


@jax.jit
def correlate(a, b):
	"""NOTE: can this be implemented as a transposition rule according to:
	https://jax.readthedocs.io/en/latest/_modules/jax/_src/lax/fft.html
	"""
	fa = jrfft(a).conj()
	fb = jrfft(b)
	return jirfft(fa * fb)


def psf_convolve(psf):
	"""statically bind psf arg"""
	psf = jax.numpy.fft.ifftshift(psf)
	return lambda a: convolve(psf, a)


def psf_correlate(psf):
	"""statically bind psf arg. psf assumed to be centered"""
	psf = jax.numpy.fft.ifftshift(psf)
	return lambda a: correlate(psf, a)


import types

def build_matrix(func, shape):
	"""explicitly evaluate coeeficient matrix of linear operator func by calling it repeatedly with delta functions"""
	i, j = shape
	Z = []
	I = jnp.eye(i, j)
	for r in range(i):
		z = func(I[r])
		Z.append(z)
	return jnp.array(Z)


func = psf_convolve(psf_jax)

arr = types.SimpleNamespace(shape=signal_jax.shape, dtype=np.float32)
func_trans = lambda a: jax.linear_transpose(func, arr)(a)[0]

N = len(signal)

plt.figure()
M = build_matrix(func, (N, N)).T
plt.imshow(M)

plt.figure()
M = build_matrix(func_trans, (N, N))
plt.imshow(M)
plt.show()
```

- [ ] If applicable, include full error messages/tracebacks.


Changed files:
jax/_src/lax/fft.py, tests/fft_test.py

diff --git a/jax/_src/lax/fft.py b/jax/_src/lax/fft.py
index 5774fdbe0..7ac0c43cd 100644
--- a/jax/_src/lax/fft.py
+++ b/jax/_src/lax/fft.py
@@ -114,7 +114,9 @@ def _irfft_transpose(t, fft_lengths):
   scale = 1 / prod(fft_lengths)
   out = scale * mask * x
   assert out.dtype == _complex_dtype(t.dtype), (out.dtype, t.dtype)
-  return out
+  # Use JAX's convention for complex gradients
+  # https://github.com/google/jax/issues/6223#issuecomment-807740707
+  return lax.conj(out)
 
 def fft_transpose_rule(t, operand, fft_type, fft_lengths):
   if fft_type == xla_client.FftType.RFFT:
diff --git a/tests/fft_test.py b/tests/fft_test.py
index d40f050c4..17f9e22d0 100644
--- a/tests/fft_test.py
+++ b/tests/fft_test.py
@@ -20,6 +20,7 @@ import numpy as np
 from absl.testing import absltest
 from absl.testing import parameterized
 
+import jax
 from jax import lax
 from jax import numpy as jnp
 from jax import test_util as jtu
@@ -116,6 +117,21 @@ class FftTest(jtu.JaxTestCase):
       tol = 0.15
       jtu.check_grads(jnp_fn, args_maker(), order=2, atol=tol, rtol=tol)
 
+  def testIrfftTranspose(self):
+    # regression test for https://github.com/google/jax/issues/6223
+    def build_matrix(linear_func, size):
+      return jax.vmap(linear_func)(jnp.eye(size, size))
+
+    def func(x):
+      return jnp.fft.irfft(jnp.concatenate([jnp.zeros(1), x[:2] + 1j*x[2:]]))
+
+    def func_transpose(x):
+      return jax.linear_transpose(func, x)(x)[0]
+
+    matrix = build_matrix(func, 4)
+    matrix2 = build_matrix(func_transpose, 4).T
+    self.assertAllClose(matrix, matrix2)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_inverse={}_real={}".format(inverse, real),
        "inverse": inverse, "real": real}