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.



Discussion:
@mattjj @hawkinsp @shoyer do any of you have cycles to look into this? (Or cc someone else who may)
I'm not quite sure what's going on yet (you may be correct about a bad transpose rule somewhere), but here's a slightly simplified example that seems to reproduce the same issue:
```python
import numpy as np
import jax
import jax.numpy as jnp

def build_matrix(func, size):
  return jax.vmap(func)(jnp.eye(size, size))

def identity(x):
  # return x  # this works
  return jnp.fft.irfft(jnp.fft.rfft(x))  # this doesn't

identity_trans = lambda a: jax.linear_transpose(identity, a)(a)[0]

N = 4
M_forward = build_matrix(identity, N)
M_backward = build_matrix(identity_trans, N).T

print('Forward matrix:')
print(M_forward)
print('Backward matrix:')
print(M_backward)
```
```
Forward matrix:
[[1. 0. 0. 0.]
 [0. 1. 0. 0.]
 [0. 0. 1. 0.]
 [0. 0. 0. 1.]]
Backward matrix:
[[1. 0. 0. 0.]
 [0. 0. 0. 1.]
 [0. 0. 1. 0.]
 [0. 1. 0. 0.]]
```
A bit more experimentation suggests that the culprit may be the transpose of `irfft`, which appears to be off by a complex conjugate:
```

import numpy as np
import jax
import jax.numpy as jnp

def build_matrix(func, size):
  return jax.vmap(func)(jnp.eye(size, size))

def func(x):
  # return jnp.fft.irfft(jnp.fft.rfft(x))
  # return jnp.fft.fft(x).real
  # return jnp.concatenate([jnp.fft.rfft(x).real, jnp.zeros(1)])
  return jnp.fft.irfft(jnp.concatenate([jnp.zeros(1), x[:2] + 1j*x[2:]]))

func_trans = lambda a: jax.linear_transpose(func, a)(a)[0]

N = 4
M_forward = build_matrix(func, N)
M_backward = build_matrix(func_trans, N).T
print('Forward matrix:')
print(M_forward)
print('Backward matrix:')
print(M_backward)
```
```
Forward matrix:
[[ 0.5   0.   -0.5   0.  ]
 [ 0.25 -0.25  0.25 -0.25]
 [ 0.   -0.5   0.    0.5 ]
 [ 0.    0.    0.    0.  ]]
Backward matrix:
[[ 0.5   0.   -0.5   0.  ]
 [ 0.25 -0.25  0.25 -0.25]
 [ 0.    0.5   0.   -0.5 ]
 [ 0.    0.    0.    0.  ]]
```

This makes sense, given that the transpose rule was copied from TensorFlow, which follows a [different convention for complex gradients](https://github.com/google/jax/issues/4891).

Resolving patch:
diff --git a/jax/_src/lax/fft.py b/jax/_src/lax/fft.py
--- 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:


Test patch:
diff --git a/tests/fft_test.py b/tests/fft_test.py
--- a/tests/fft_test.py
+++ b/tests/fft_test.py
@@ -20,6 +20,7 @@
 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 @@ def testFftn(self, inverse, real, shape, dtype, axes):
       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}