We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
fft2


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `ivy/functional/frontends/jax/numpy/fft.py`
Content:
```
1 # local
2 import ivy
3 from ivy.functional.frontends.jax.func_wrapper import to_ivy_arrays_and_back
4 from ivy.func_wrapper import with_unsupported_dtypes
5 
6 
7 @to_ivy_arrays_and_back
8 def fft(a, n=None, axis=-1, norm=None):
9     if norm is None:
10         norm = "backward"
11     return ivy.fft(a, axis, norm=norm, n=n)
12 
13 
14 @to_ivy_arrays_and_back
15 @with_unsupported_dtypes({"2.4.2 and below": ("float16", "bfloat16")}, "paddle")
16 def fftshift(x, axes=None, name=None):
17     shape = x.shape
18 
19     if axes is None:
20         axes = tuple(range(x.ndim))
21         shifts = [(dim // 2) for dim in shape]
22     elif isinstance(axes, int):
23         shifts = shape[axes] // 2
24     else:
25         shifts = [shape[ax] // 2 for ax in axes]
26 
27     roll = ivy.roll(x, shifts, axis=axes)
28 
29     return roll
30
```
--- 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/ivy/functional/frontends/jax/numpy/fft.py b/ivy/functional/frontends/jax/numpy/fft.py
--- a/ivy/functional/frontends/jax/numpy/fft.py
+++ b/ivy/functional/frontends/jax/numpy/fft.py
@@ -11,6 +11,13 @@
     return ivy.fft(a, axis, norm=norm, n=n)
 
 
+@to_ivy_arrays_and_back
+def fft2(a, s=None, axes=(-2, -1), norm=None):
+    if norm is None:
+        norm = "backward"
+    return ivy.array(ivy.fft2(a, s=s, dim=axes, norm=norm), dtype=ivy.dtype(a))
+
+
 @to_ivy_arrays_and_back
 @with_unsupported_dtypes({"2.4.2 and below": ("float16", "bfloat16")}, "paddle")
 def fftshift(x, axes=None, name=None):

{"golden_diff": "diff --git a/ivy/functional/frontends/jax/numpy/fft.py b/ivy/functional/frontends/jax/numpy/fft.py\n--- a/ivy/functional/frontends/jax/numpy/fft.py\n+++ b/ivy/functional/frontends/jax/numpy/fft.py\n@@ -11,6 +11,13 @@\n     return ivy.fft(a, axis, norm=norm, n=n)\n \n \n+@to_ivy_arrays_and_back\n+def fft2(a, s=None, axes=(-2, -1), norm=None):\n+    if norm is None:\n+        norm = \"backward\"\n+    return ivy.array(ivy.fft2(a, s=s, dim=axes, norm=norm), dtype=ivy.dtype(a))\n+\n+\n @to_ivy_arrays_and_back\n @with_unsupported_dtypes({\"2.4.2 and below\": (\"float16\", \"bfloat16\")}, \"paddle\")\n def fftshift(x, axes=None, name=None):\n", "issue": "fft2\n\n", "before_files": [{"content": "# local\nimport ivy\nfrom ivy.functional.frontends.jax.func_wrapper import to_ivy_arrays_and_back\nfrom ivy.func_wrapper import with_unsupported_dtypes\n\n\n@to_ivy_arrays_and_back\ndef fft(a, n=None, axis=-1, norm=None):\n    if norm is None:\n        norm = \"backward\"\n    return ivy.fft(a, axis, norm=norm, n=n)\n\n\n@to_ivy_arrays_and_back\n@with_unsupported_dtypes({\"2.4.2 and below\": (\"float16\", \"bfloat16\")}, \"paddle\")\ndef fftshift(x, axes=None, name=None):\n    shape = x.shape\n\n    if axes is None:\n        axes = tuple(range(x.ndim))\n        shifts = [(dim // 2) for dim in shape]\n    elif isinstance(axes, int):\n        shifts = shape[axes] // 2\n    else:\n        shifts = [shape[ax] // 2 for ax in axes]\n\n    roll = ivy.roll(x, shifts, axis=axes)\n\n    return roll\n", "path": "ivy/functional/frontends/jax/numpy/fft.py"}], "after_files": [{"content": "# local\nimport ivy\nfrom ivy.functional.frontends.jax.func_wrapper import to_ivy_arrays_and_back\nfrom ivy.func_wrapper import with_unsupported_dtypes\n\n\n@to_ivy_arrays_and_back\ndef fft(a, n=None, axis=-1, norm=None):\n    if norm is None:\n        norm = \"backward\"\n    return ivy.fft(a, axis, norm=norm, n=n)\n\n\n@to_ivy_arrays_and_back\ndef fft2(a, s=None, axes=(-2, -1), norm=None):\n    if norm is None:\n        norm = \"backward\"\n    return ivy.array(ivy.fft2(a, s=s, dim=axes, norm=norm), dtype=ivy.dtype(a))\n\n\n@to_ivy_arrays_and_back\n@with_unsupported_dtypes({\"2.4.2 and below\": (\"float16\", \"bfloat16\")}, \"paddle\")\ndef fftshift(x, axes=None, name=None):\n    shape = x.shape\n\n    if axes is None:\n        axes = tuple(range(x.ndim))\n        shifts = [(dim // 2) for dim in shape]\n    elif isinstance(axes, int):\n        shifts = shape[axes] // 2\n    else:\n        shifts = [shape[ax] // 2 for ax in axes]\n\n    roll = ivy.roll(x, shifts, axis=axes)\n\n    return roll\n", "path": "ivy/functional/frontends/jax/numpy/fft.py"}]}