We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `jax/_src/lax/fft.py`
Content:
```
1 # Copyright 2019 Google LLC
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 
16 from functools import partial
17 
18 import numpy as np
19 
20 from jax.api import jit, linear_transpose, ShapeDtypeStruct
21 from jax.core import Primitive
22 from jax.interpreters import xla
23 from jax._src.util import prod
24 from jax import dtypes, lax
25 from jax.lib import xla_client
26 from jax.interpreters import ad
27 from jax.interpreters import batching
28 from jax.lib import pocketfft
29 
30 xops = xla_client.ops
31 
32 __all__ = [
33   "fft",
34   "fft_p",
35 ]
36 
37 def _promote_to_complex(arg):
38   dtype = dtypes.result_type(arg, np.complex64)
39   return lax.convert_element_type(arg, dtype)
40 
41 def _promote_to_real(arg):
42   dtype = dtypes.result_type(arg, np.float32)
43   return lax.convert_element_type(arg, dtype)
44 
45 def fft(x, fft_type, fft_lengths):
46   if fft_type == xla_client.FftType.RFFT:
47     if np.iscomplexobj(x):
48       raise ValueError("only real valued inputs supported for rfft")
49     x = _promote_to_real(x)
50   else:
51     x = _promote_to_complex(x)
52   if len(fft_lengths) == 0:
53     # XLA FFT doesn't support 0-rank.
54     return x
55   fft_lengths = tuple(fft_lengths)
56   return fft_p.bind(x, fft_type=fft_type, fft_lengths=fft_lengths)
57 
58 def fft_impl(x, fft_type, fft_lengths):
59   return xla.apply_primitive(fft_p, x, fft_type=fft_type, fft_lengths=fft_lengths)
60 
61 _complex_dtype = lambda dtype: (np.zeros((), dtype) + np.zeros((), np.complex64)).dtype
62 _real_dtype = lambda dtype: np.zeros((), dtype).real.dtype
63 _is_even = lambda x: x % 2 == 0
64 
65 def fft_abstract_eval(x, fft_type, fft_lengths):
66   if fft_type == xla_client.FftType.RFFT:
67     shape = (x.shape[:-len(fft_lengths)] + fft_lengths[:-1]
68              + (fft_lengths[-1] // 2 + 1,))
69     dtype = _complex_dtype(x.dtype)
70   elif fft_type == xla_client.FftType.IRFFT:
71     shape = x.shape[:-len(fft_lengths)] + fft_lengths
72     dtype = _real_dtype(x.dtype)
73   else:
74     shape = x.shape
75     dtype = x.dtype
76   return x.update(shape=shape, dtype=dtype)
77 
78 def fft_translation_rule(c, x, fft_type, fft_lengths):
79   return xops.Fft(x, fft_type, fft_lengths)
80 
81 def _naive_rfft(x, fft_lengths):
82   y = fft(x, xla_client.FftType.FFT, fft_lengths)
83   n = fft_lengths[-1]
84   return y[..., : n//2 + 1]
85 
86 @partial(jit, static_argnums=1)
87 def _rfft_transpose(t, fft_lengths):
88   # The transpose of RFFT can't be expressed only in terms of irfft. Instead of
89   # manually building up larger twiddle matrices (which would increase the
90   # asymptotic complexity and is also rather complicated), we rely JAX to
91   # transpose a naive RFFT implementation.
92   dummy_shape = t.shape[:-len(fft_lengths)] + fft_lengths
93   dummy_primal = ShapeDtypeStruct(dummy_shape, _real_dtype(t.dtype))
94   transpose = linear_transpose(
95       partial(_naive_rfft, fft_lengths=fft_lengths), dummy_primal)
96   result, = transpose(t)
97   assert result.dtype == _real_dtype(t.dtype), (result.dtype, t.dtype)
98   return result
99 
100 def _irfft_transpose(t, fft_lengths):
101   # The transpose of IRFFT is the RFFT of the cotangent times a scaling
102   # factor and a mask. The mask scales the cotangent for the Hermitian
103   # symmetric components of the RFFT by a factor of two, since these components
104   # are de-duplicated in the RFFT.
105   x = fft(t, xla_client.FftType.RFFT, fft_lengths)
106   n = x.shape[-1]
107   is_odd = fft_lengths[-1] % 2
108   full = partial(lax.full_like, t, dtype=t.dtype)
109   mask = lax.concatenate(
110       [full(1.0, shape=(1,)),
111        full(2.0, shape=(n - 2 + is_odd,)),
112        full(1.0, shape=(1 - is_odd,))],
113       dimension=0)
114   scale = 1 / prod(fft_lengths)
115   out = scale * mask * x
116   assert out.dtype == _complex_dtype(t.dtype), (out.dtype, t.dtype)
117   return out
118 
119 def fft_transpose_rule(t, operand, fft_type, fft_lengths):
120   if fft_type == xla_client.FftType.RFFT:
121     result = _rfft_transpose(t, fft_lengths)
122   elif fft_type == xla_client.FftType.IRFFT:
123     result = _irfft_transpose(t, fft_lengths)
124   else:
125     result = fft(t, fft_type, fft_lengths)
126   return result,
127 
128 def fft_batching_rule(batched_args, batch_dims, fft_type, fft_lengths):
129   x, = batched_args
130   bd, = batch_dims
131   x = batching.moveaxis(x, bd, 0)
132   return fft(x, fft_type, fft_lengths), 0
133 
134 fft_p = Primitive('fft')
135 fft_p.def_impl(fft_impl)
136 fft_p.def_abstract_eval(fft_abstract_eval)
137 xla.translations[fft_p] = fft_translation_rule
138 ad.deflinear2(fft_p, fft_transpose_rule)
139 batching.primitive_batchers[fft_p] = fft_batching_rule
140 if pocketfft:
141   xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft
142
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
   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:

{"golden_diff": "diff --git a/jax/_src/lax/fft.py b/jax/_src/lax/fft.py\n--- a/jax/_src/lax/fft.py\n+++ b/jax/_src/lax/fft.py\n@@ -114,7 +114,9 @@\n   scale = 1 / prod(fft_lengths)\n   out = scale * mask * x\n   assert out.dtype == _complex_dtype(t.dtype), (out.dtype, t.dtype)\n-  return out\n+  # Use JAX's convention for complex gradients\n+  # https://github.com/google/jax/issues/6223#issuecomment-807740707\n+  return lax.conj(out)\n \n def fft_transpose_rule(t, operand, fft_type, fft_lengths):\n   if fft_type == xla_client.FftType.RFFT:\n", "issue": "linear_transpose involving fft seemingly incorrect\nI 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.\r\n\r\nI 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).\r\n\r\nSeems to me like a bug in the implementation of the transpose rules for ffts; but again not qualified myself to spot it.\r\n\r\nWhile 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)?\r\n\r\n\r\nCode to reproduce:\r\n\r\n```python\r\nimport numpy as np\r\n\r\nimport jax\r\nfrom jax import numpy as jnp\r\nimport matplotlib.pyplot as plt\r\n\r\n\r\nnp.random.seed(0)\r\n\r\nsignal = np.cumsum(np.random.randn(2**8))\r\nsignal_jax = jnp.array(signal)\r\n\r\nx = np.linspace(-1, 1, len(signal))\r\npsf = np.clip(0.2 - np.abs(x), 0, 1) * (x > 0)\r\npsf /= psf.sum()\r\npsf_jax = jnp.array(psf)\r\n\r\njrfft = jax.jit(jnp.fft.rfft)\r\njirfft = jax.jit(jnp.fft.irfft)\r\n\r\n\r\n@jax.jit\r\ndef convolve(a, b):\r\n\tfa = jrfft(a)\r\n\tfb = jrfft(b)\r\n\treturn jirfft(fa * fb)\r\n\r\n\r\n@jax.jit\r\ndef correlate(a, b):\r\n\t\"\"\"NOTE: can this be implemented as a transposition rule according to:\r\n\thttps://jax.readthedocs.io/en/latest/_modules/jax/_src/lax/fft.html\r\n\t\"\"\"\r\n\tfa = jrfft(a).conj()\r\n\tfb = jrfft(b)\r\n\treturn jirfft(fa * fb)\r\n\r\n\r\ndef psf_convolve(psf):\r\n\t\"\"\"statically bind psf arg\"\"\"\r\n\tpsf = jax.numpy.fft.ifftshift(psf)\r\n\treturn lambda a: convolve(psf, a)\r\n\r\n\r\ndef psf_correlate(psf):\r\n\t\"\"\"statically bind psf arg. psf assumed to be centered\"\"\"\r\n\tpsf = jax.numpy.fft.ifftshift(psf)\r\n\treturn lambda a: correlate(psf, a)\r\n\r\n\r\nimport types\r\n\r\ndef build_matrix(func, shape):\r\n\t\"\"\"explicitly evaluate coeeficient matrix of linear operator func by calling it repeatedly with delta functions\"\"\"\r\n\ti, j = shape\r\n\tZ = []\r\n\tI = jnp.eye(i, j)\r\n\tfor r in range(i):\r\n\t\tz = func(I[r])\r\n\t\tZ.append(z)\r\n\treturn jnp.array(Z)\r\n\r\n\r\nfunc = psf_convolve(psf_jax)\r\n\r\narr = types.SimpleNamespace(shape=signal_jax.shape, dtype=np.float32)\r\nfunc_trans = lambda a: jax.linear_transpose(func, arr)(a)[0]\r\n\r\nN = len(signal)\r\n\r\nplt.figure()\r\nM = build_matrix(func, (N, N)).T\r\nplt.imshow(M)\r\n\r\nplt.figure()\r\nM = build_matrix(func_trans, (N, N))\r\nplt.imshow(M)\r\nplt.show()\r\n```\r\n\r\n- [ ] If applicable, include full error messages/tracebacks.\r\n\n", "before_files": [{"content": "# Copyright 2019 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom functools import partial\n\nimport numpy as np\n\nfrom jax.api import jit, linear_transpose, ShapeDtypeStruct\nfrom jax.core import Primitive\nfrom jax.interpreters import xla\nfrom jax._src.util import prod\nfrom jax import dtypes, lax\nfrom jax.lib import xla_client\nfrom jax.interpreters import ad\nfrom jax.interpreters import batching\nfrom jax.lib import pocketfft\n\nxops = xla_client.ops\n\n__all__ = [\n  \"fft\",\n  \"fft_p\",\n]\n\ndef _promote_to_complex(arg):\n  dtype = dtypes.result_type(arg, np.complex64)\n  return lax.convert_element_type(arg, dtype)\n\ndef _promote_to_real(arg):\n  dtype = dtypes.result_type(arg, np.float32)\n  return lax.convert_element_type(arg, dtype)\n\ndef fft(x, fft_type, fft_lengths):\n  if fft_type == xla_client.FftType.RFFT:\n    if np.iscomplexobj(x):\n      raise ValueError(\"only real valued inputs supported for rfft\")\n    x = _promote_to_real(x)\n  else:\n    x = _promote_to_complex(x)\n  if len(fft_lengths) == 0:\n    # XLA FFT doesn't support 0-rank.\n    return x\n  fft_lengths = tuple(fft_lengths)\n  return fft_p.bind(x, fft_type=fft_type, fft_lengths=fft_lengths)\n\ndef fft_impl(x, fft_type, fft_lengths):\n  return xla.apply_primitive(fft_p, x, fft_type=fft_type, fft_lengths=fft_lengths)\n\n_complex_dtype = lambda dtype: (np.zeros((), dtype) + np.zeros((), np.complex64)).dtype\n_real_dtype = lambda dtype: np.zeros((), dtype).real.dtype\n_is_even = lambda x: x % 2 == 0\n\ndef fft_abstract_eval(x, fft_type, fft_lengths):\n  if fft_type == xla_client.FftType.RFFT:\n    shape = (x.shape[:-len(fft_lengths)] + fft_lengths[:-1]\n             + (fft_lengths[-1] // 2 + 1,))\n    dtype = _complex_dtype(x.dtype)\n  elif fft_type == xla_client.FftType.IRFFT:\n    shape = x.shape[:-len(fft_lengths)] + fft_lengths\n    dtype = _real_dtype(x.dtype)\n  else:\n    shape = x.shape\n    dtype = x.dtype\n  return x.update(shape=shape, dtype=dtype)\n\ndef fft_translation_rule(c, x, fft_type, fft_lengths):\n  return xops.Fft(x, fft_type, fft_lengths)\n\ndef _naive_rfft(x, fft_lengths):\n  y = fft(x, xla_client.FftType.FFT, fft_lengths)\n  n = fft_lengths[-1]\n  return y[..., : n//2 + 1]\n\n@partial(jit, static_argnums=1)\ndef _rfft_transpose(t, fft_lengths):\n  # The transpose of RFFT can't be expressed only in terms of irfft. Instead of\n  # manually building up larger twiddle matrices (which would increase the\n  # asymptotic complexity and is also rather complicated), we rely JAX to\n  # transpose a naive RFFT implementation.\n  dummy_shape = t.shape[:-len(fft_lengths)] + fft_lengths\n  dummy_primal = ShapeDtypeStruct(dummy_shape, _real_dtype(t.dtype))\n  transpose = linear_transpose(\n      partial(_naive_rfft, fft_lengths=fft_lengths), dummy_primal)\n  result, = transpose(t)\n  assert result.dtype == _real_dtype(t.dtype), (result.dtype, t.dtype)\n  return result\n\ndef _irfft_transpose(t, fft_lengths):\n  # The transpose of IRFFT is the RFFT of the cotangent times a scaling\n  # factor and a mask. The mask scales the cotangent for the Hermitian\n  # symmetric components of the RFFT by a factor of two, since these components\n  # are de-duplicated in the RFFT.\n  x = fft(t, xla_client.FftType.RFFT, fft_lengths)\n  n = x.shape[-1]\n  is_odd = fft_lengths[-1] % 2\n  full = partial(lax.full_like, t, dtype=t.dtype)\n  mask = lax.concatenate(\n      [full(1.0, shape=(1,)),\n       full(2.0, shape=(n - 2 + is_odd,)),\n       full(1.0, shape=(1 - is_odd,))],\n      dimension=0)\n  scale = 1 / prod(fft_lengths)\n  out = scale * mask * x\n  assert out.dtype == _complex_dtype(t.dtype), (out.dtype, t.dtype)\n  return out\n\ndef fft_transpose_rule(t, operand, fft_type, fft_lengths):\n  if fft_type == xla_client.FftType.RFFT:\n    result = _rfft_transpose(t, fft_lengths)\n  elif fft_type == xla_client.FftType.IRFFT:\n    result = _irfft_transpose(t, fft_lengths)\n  else:\n    result = fft(t, fft_type, fft_lengths)\n  return result,\n\ndef fft_batching_rule(batched_args, batch_dims, fft_type, fft_lengths):\n  x, = batched_args\n  bd, = batch_dims\n  x = batching.moveaxis(x, bd, 0)\n  return fft(x, fft_type, fft_lengths), 0\n\nfft_p = Primitive('fft')\nfft_p.def_impl(fft_impl)\nfft_p.def_abstract_eval(fft_abstract_eval)\nxla.translations[fft_p] = fft_translation_rule\nad.deflinear2(fft_p, fft_transpose_rule)\nbatching.primitive_batchers[fft_p] = fft_batching_rule\nif pocketfft:\n  xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft\n", "path": "jax/_src/lax/fft.py"}], "after_files": [{"content": "# Copyright 2019 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom functools import partial\n\nimport numpy as np\n\nfrom jax.api import jit, linear_transpose, ShapeDtypeStruct\nfrom jax.core import Primitive\nfrom jax.interpreters import xla\nfrom jax._src.util import prod\nfrom jax import dtypes, lax\nfrom jax.lib import xla_client\nfrom jax.interpreters import ad\nfrom jax.interpreters import batching\nfrom jax.lib import pocketfft\n\nxops = xla_client.ops\n\n__all__ = [\n  \"fft\",\n  \"fft_p\",\n]\n\ndef _promote_to_complex(arg):\n  dtype = dtypes.result_type(arg, np.complex64)\n  return lax.convert_element_type(arg, dtype)\n\ndef _promote_to_real(arg):\n  dtype = dtypes.result_type(arg, np.float32)\n  return lax.convert_element_type(arg, dtype)\n\ndef fft(x, fft_type, fft_lengths):\n  if fft_type == xla_client.FftType.RFFT:\n    if np.iscomplexobj(x):\n      raise ValueError(\"only real valued inputs supported for rfft\")\n    x = _promote_to_real(x)\n  else:\n    x = _promote_to_complex(x)\n  if len(fft_lengths) == 0:\n    # XLA FFT doesn't support 0-rank.\n    return x\n  fft_lengths = tuple(fft_lengths)\n  return fft_p.bind(x, fft_type=fft_type, fft_lengths=fft_lengths)\n\ndef fft_impl(x, fft_type, fft_lengths):\n  return xla.apply_primitive(fft_p, x, fft_type=fft_type, fft_lengths=fft_lengths)\n\n_complex_dtype = lambda dtype: (np.zeros((), dtype) + np.zeros((), np.complex64)).dtype\n_real_dtype = lambda dtype: np.zeros((), dtype).real.dtype\n_is_even = lambda x: x % 2 == 0\n\ndef fft_abstract_eval(x, fft_type, fft_lengths):\n  if fft_type == xla_client.FftType.RFFT:\n    shape = (x.shape[:-len(fft_lengths)] + fft_lengths[:-1]\n             + (fft_lengths[-1] // 2 + 1,))\n    dtype = _complex_dtype(x.dtype)\n  elif fft_type == xla_client.FftType.IRFFT:\n    shape = x.shape[:-len(fft_lengths)] + fft_lengths\n    dtype = _real_dtype(x.dtype)\n  else:\n    shape = x.shape\n    dtype = x.dtype\n  return x.update(shape=shape, dtype=dtype)\n\ndef fft_translation_rule(c, x, fft_type, fft_lengths):\n  return xops.Fft(x, fft_type, fft_lengths)\n\ndef _naive_rfft(x, fft_lengths):\n  y = fft(x, xla_client.FftType.FFT, fft_lengths)\n  n = fft_lengths[-1]\n  return y[..., : n//2 + 1]\n\n@partial(jit, static_argnums=1)\ndef _rfft_transpose(t, fft_lengths):\n  # The transpose of RFFT can't be expressed only in terms of irfft. Instead of\n  # manually building up larger twiddle matrices (which would increase the\n  # asymptotic complexity and is also rather complicated), we rely JAX to\n  # transpose a naive RFFT implementation.\n  dummy_shape = t.shape[:-len(fft_lengths)] + fft_lengths\n  dummy_primal = ShapeDtypeStruct(dummy_shape, _real_dtype(t.dtype))\n  transpose = linear_transpose(\n      partial(_naive_rfft, fft_lengths=fft_lengths), dummy_primal)\n  result, = transpose(t)\n  assert result.dtype == _real_dtype(t.dtype), (result.dtype, t.dtype)\n  return result\n\ndef _irfft_transpose(t, fft_lengths):\n  # The transpose of IRFFT is the RFFT of the cotangent times a scaling\n  # factor and a mask. The mask scales the cotangent for the Hermitian\n  # symmetric components of the RFFT by a factor of two, since these components\n  # are de-duplicated in the RFFT.\n  x = fft(t, xla_client.FftType.RFFT, fft_lengths)\n  n = x.shape[-1]\n  is_odd = fft_lengths[-1] % 2\n  full = partial(lax.full_like, t, dtype=t.dtype)\n  mask = lax.concatenate(\n      [full(1.0, shape=(1,)),\n       full(2.0, shape=(n - 2 + is_odd,)),\n       full(1.0, shape=(1 - is_odd,))],\n      dimension=0)\n  scale = 1 / prod(fft_lengths)\n  out = scale * mask * x\n  assert out.dtype == _complex_dtype(t.dtype), (out.dtype, t.dtype)\n  # Use JAX's convention for complex gradients\n  # https://github.com/google/jax/issues/6223#issuecomment-807740707\n  return lax.conj(out)\n\ndef fft_transpose_rule(t, operand, fft_type, fft_lengths):\n  if fft_type == xla_client.FftType.RFFT:\n    result = _rfft_transpose(t, fft_lengths)\n  elif fft_type == xla_client.FftType.IRFFT:\n    result = _irfft_transpose(t, fft_lengths)\n  else:\n    result = fft(t, fft_type, fft_lengths)\n  return result,\n\ndef fft_batching_rule(batched_args, batch_dims, fft_type, fft_lengths):\n  x, = batched_args\n  bd, = batch_dims\n  x = batching.moveaxis(x, bd, 0)\n  return fft(x, fft_type, fft_lengths), 0\n\nfft_p = Primitive('fft')\nfft_p.def_impl(fft_impl)\nfft_p.def_abstract_eval(fft_abstract_eval)\nxla.translations[fft_p] = fft_translation_rule\nad.deflinear2(fft_p, fft_transpose_rule)\nbatching.primitive_batchers[fft_p] = fft_batching_rule\nif pocketfft:\n  xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft\n", "path": "jax/_src/lax/fft.py"}]}