We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
JTFS needs a JAX frontend
v0.4.0

Already addressed by #1001 (@cyrusvahidi)

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `kymatio/scattering1d/frontend/jax_frontend.py`
Content:
```
1 from ...frontend.jax_frontend import ScatteringJax
2 from .numpy_frontend import ScatteringNumPy1D
3 from .base_frontend import ScatteringBase1D
4 
5 class ScatteringJax1D(ScatteringJax, ScatteringNumPy1D):
6     # This class inherits the attribute "frontend" from ScatteringJax
7     # It overrides the __init__ function present in ScatteringNumPy1D
8     # in order to add the default argument for backend and call the
9     # ScatteringJax.__init__
10     # Through ScatteringBase1D._instantiate_backend the jax backend will
11     # be loaded
12 
13 
14     def __init__(self, J, shape, Q=1, T=None, stride=None, max_order=2,
15             oversampling=0, out_type='array', backend='jax'):
16 
17         ScatteringJax.__init__(self)
18         ScatteringBase1D.__init__(self, J, shape, Q, T, stride, max_order,
19                 oversampling, out_type, backend)
20         ScatteringBase1D._instantiate_backend(self, 'kymatio.scattering1d.backend.')
21         ScatteringBase1D.build(self)
22         ScatteringBase1D.create_filters(self)
23 
24 ScatteringJax1D._document()
25
```
--- 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/kymatio/scattering1d/frontend/jax_frontend.py b/kymatio/scattering1d/frontend/jax_frontend.py
--- a/kymatio/scattering1d/frontend/jax_frontend.py
+++ b/kymatio/scattering1d/frontend/jax_frontend.py
@@ -1,6 +1,7 @@
 from ...frontend.jax_frontend import ScatteringJax
-from .numpy_frontend import ScatteringNumPy1D
-from .base_frontend import ScatteringBase1D
+from .numpy_frontend import ScatteringNumPy1D, TimeFrequencyScatteringNumPy
+from .base_frontend import ScatteringBase1D, TimeFrequencyScatteringBase
+
 
 class ScatteringJax1D(ScatteringJax, ScatteringNumPy1D):
     # This class inherits the attribute "frontend" from ScatteringJax
@@ -10,15 +11,77 @@
     # Through ScatteringBase1D._instantiate_backend the jax backend will
     # be loaded
 
-
-    def __init__(self, J, shape, Q=1, T=None, stride=None, max_order=2,
-            oversampling=0, out_type='array', backend='jax'):
+    def __init__(
+        self,
+        J,
+        shape,
+        Q=1,
+        T=None,
+        stride=None,
+        max_order=2,
+        oversampling=0,
+        out_type="array",
+        backend="jax",
+    ):
 
         ScatteringJax.__init__(self)
-        ScatteringBase1D.__init__(self, J, shape, Q, T, stride, max_order,
-                oversampling, out_type, backend)
-        ScatteringBase1D._instantiate_backend(self, 'kymatio.scattering1d.backend.')
+        ScatteringBase1D.__init__(
+            self, J, shape, Q, T, stride, max_order, oversampling, out_type, backend
+        )
+        ScatteringBase1D._instantiate_backend(self, "kymatio.scattering1d.backend.")
         ScatteringBase1D.build(self)
         ScatteringBase1D.create_filters(self)
 
+
 ScatteringJax1D._document()
+
+
+class TimeFrequencyScatteringJax(ScatteringJax, TimeFrequencyScatteringNumPy):
+    # This class inherits the attribute "frontend" from ScatteringJax
+    # It overrides the __init__ function present in TimeFrequencyScatteringNumPy
+    # in order to add the default argument for backend and call the
+    # ScatteringJax.__init__
+    # Through TimeFrequencyScatteringBase._instantiate_backend the jax backend will
+    # be loaded
+
+    def __init__(
+        self,
+        *,
+        J,
+        J_fr,
+        shape,
+        Q,
+        T=None,
+        stride=None,
+        Q_fr=1,
+        F=None,
+        stride_fr=None,
+        out_type="array",
+        format="joint",
+        backend="jax"
+    ):
+
+        ScatteringJax.__init__(self)
+        TimeFrequencyScatteringBase.__init__(
+            self,
+            J=J,
+            J_fr=J_fr,
+            shape=shape,
+            Q=Q,
+            T=T,
+            stride=stride,
+            Q_fr=Q_fr,
+            F=F,
+            stride_fr=stride_fr,
+            out_type=out_type,
+            format=format,
+            backend=backend,
+        )
+        ScatteringBase1D._instantiate_backend(self, "kymatio.scattering1d.backend.")
+        TimeFrequencyScatteringBase.build(self)
+        TimeFrequencyScatteringBase.create_filters(self)
+
+
+TimeFrequencyScatteringJax._document()
+
+__all__ = ["ScatteringJax1D", "TimeFrequencyScatteringJax"]

{"golden_diff": "diff --git a/kymatio/scattering1d/frontend/jax_frontend.py b/kymatio/scattering1d/frontend/jax_frontend.py\n--- a/kymatio/scattering1d/frontend/jax_frontend.py\n+++ b/kymatio/scattering1d/frontend/jax_frontend.py\n@@ -1,6 +1,7 @@\n from ...frontend.jax_frontend import ScatteringJax\n-from .numpy_frontend import ScatteringNumPy1D\n-from .base_frontend import ScatteringBase1D\n+from .numpy_frontend import ScatteringNumPy1D, TimeFrequencyScatteringNumPy\n+from .base_frontend import ScatteringBase1D, TimeFrequencyScatteringBase\n+\n \n class ScatteringJax1D(ScatteringJax, ScatteringNumPy1D):\n     # This class inherits the attribute \"frontend\" from ScatteringJax\n@@ -10,15 +11,77 @@\n     # Through ScatteringBase1D._instantiate_backend the jax backend will\n     # be loaded\n \n-\n-    def __init__(self, J, shape, Q=1, T=None, stride=None, max_order=2,\n-            oversampling=0, out_type='array', backend='jax'):\n+    def __init__(\n+        self,\n+        J,\n+        shape,\n+        Q=1,\n+        T=None,\n+        stride=None,\n+        max_order=2,\n+        oversampling=0,\n+        out_type=\"array\",\n+        backend=\"jax\",\n+    ):\n \n         ScatteringJax.__init__(self)\n-        ScatteringBase1D.__init__(self, J, shape, Q, T, stride, max_order,\n-                oversampling, out_type, backend)\n-        ScatteringBase1D._instantiate_backend(self, 'kymatio.scattering1d.backend.')\n+        ScatteringBase1D.__init__(\n+            self, J, shape, Q, T, stride, max_order, oversampling, out_type, backend\n+        )\n+        ScatteringBase1D._instantiate_backend(self, \"kymatio.scattering1d.backend.\")\n         ScatteringBase1D.build(self)\n         ScatteringBase1D.create_filters(self)\n \n+\n ScatteringJax1D._document()\n+\n+\n+class TimeFrequencyScatteringJax(ScatteringJax, TimeFrequencyScatteringNumPy):\n+    # This class inherits the attribute \"frontend\" from ScatteringJax\n+    # It overrides the __init__ function present in TimeFrequencyScatteringNumPy\n+    # in order to add the default argument for backend and call the\n+    # ScatteringJax.__init__\n+    # Through TimeFrequencyScatteringBase._instantiate_backend the jax backend will\n+    # be loaded\n+\n+    def __init__(\n+        self,\n+        *,\n+        J,\n+        J_fr,\n+        shape,\n+        Q,\n+        T=None,\n+        stride=None,\n+        Q_fr=1,\n+        F=None,\n+        stride_fr=None,\n+        out_type=\"array\",\n+        format=\"joint\",\n+        backend=\"jax\"\n+    ):\n+\n+        ScatteringJax.__init__(self)\n+        TimeFrequencyScatteringBase.__init__(\n+            self,\n+            J=J,\n+            J_fr=J_fr,\n+            shape=shape,\n+            Q=Q,\n+            T=T,\n+            stride=stride,\n+            Q_fr=Q_fr,\n+            F=F,\n+            stride_fr=stride_fr,\n+            out_type=out_type,\n+            format=format,\n+            backend=backend,\n+        )\n+        ScatteringBase1D._instantiate_backend(self, \"kymatio.scattering1d.backend.\")\n+        TimeFrequencyScatteringBase.build(self)\n+        TimeFrequencyScatteringBase.create_filters(self)\n+\n+\n+TimeFrequencyScatteringJax._document()\n+\n+__all__ = [\"ScatteringJax1D\", \"TimeFrequencyScatteringJax\"]\n", "issue": "JTFS needs a JAX frontend\nv0.4.0\r\n\r\nAlready addressed by #1001 (@cyrusvahidi)\n", "before_files": [{"content": "from ...frontend.jax_frontend import ScatteringJax\nfrom .numpy_frontend import ScatteringNumPy1D\nfrom .base_frontend import ScatteringBase1D\n\nclass ScatteringJax1D(ScatteringJax, ScatteringNumPy1D):\n    # This class inherits the attribute \"frontend\" from ScatteringJax\n    # It overrides the __init__ function present in ScatteringNumPy1D\n    # in order to add the default argument for backend and call the\n    # ScatteringJax.__init__\n    # Through ScatteringBase1D._instantiate_backend the jax backend will\n    # be loaded\n\n\n    def __init__(self, J, shape, Q=1, T=None, stride=None, max_order=2,\n            oversampling=0, out_type='array', backend='jax'):\n\n        ScatteringJax.__init__(self)\n        ScatteringBase1D.__init__(self, J, shape, Q, T, stride, max_order,\n                oversampling, out_type, backend)\n        ScatteringBase1D._instantiate_backend(self, 'kymatio.scattering1d.backend.')\n        ScatteringBase1D.build(self)\n        ScatteringBase1D.create_filters(self)\n\nScatteringJax1D._document()\n", "path": "kymatio/scattering1d/frontend/jax_frontend.py"}], "after_files": [{"content": "from ...frontend.jax_frontend import ScatteringJax\nfrom .numpy_frontend import ScatteringNumPy1D, TimeFrequencyScatteringNumPy\nfrom .base_frontend import ScatteringBase1D, TimeFrequencyScatteringBase\n\n\nclass ScatteringJax1D(ScatteringJax, ScatteringNumPy1D):\n    # This class inherits the attribute \"frontend\" from ScatteringJax\n    # It overrides the __init__ function present in ScatteringNumPy1D\n    # in order to add the default argument for backend and call the\n    # ScatteringJax.__init__\n    # Through ScatteringBase1D._instantiate_backend the jax backend will\n    # be loaded\n\n    def __init__(\n        self,\n        J,\n        shape,\n        Q=1,\n        T=None,\n        stride=None,\n        max_order=2,\n        oversampling=0,\n        out_type=\"array\",\n        backend=\"jax\",\n    ):\n\n        ScatteringJax.__init__(self)\n        ScatteringBase1D.__init__(\n            self, J, shape, Q, T, stride, max_order, oversampling, out_type, backend\n        )\n        ScatteringBase1D._instantiate_backend(self, \"kymatio.scattering1d.backend.\")\n        ScatteringBase1D.build(self)\n        ScatteringBase1D.create_filters(self)\n\n\nScatteringJax1D._document()\n\n\nclass TimeFrequencyScatteringJax(ScatteringJax, TimeFrequencyScatteringNumPy):\n    # This class inherits the attribute \"frontend\" from ScatteringJax\n    # It overrides the __init__ function present in TimeFrequencyScatteringNumPy\n    # in order to add the default argument for backend and call the\n    # ScatteringJax.__init__\n    # Through TimeFrequencyScatteringBase._instantiate_backend the jax backend will\n    # be loaded\n\n    def __init__(\n        self,\n        *,\n        J,\n        J_fr,\n        shape,\n        Q,\n        T=None,\n        stride=None,\n        Q_fr=1,\n        F=None,\n        stride_fr=None,\n        out_type=\"array\",\n        format=\"joint\",\n        backend=\"jax\"\n    ):\n\n        ScatteringJax.__init__(self)\n        TimeFrequencyScatteringBase.__init__(\n            self,\n            J=J,\n            J_fr=J_fr,\n            shape=shape,\n            Q=Q,\n            T=T,\n            stride=stride,\n            Q_fr=Q_fr,\n            F=F,\n            stride_fr=stride_fr,\n            out_type=out_type,\n            format=format,\n            backend=backend,\n        )\n        ScatteringBase1D._instantiate_backend(self, \"kymatio.scattering1d.backend.\")\n        TimeFrequencyScatteringBase.build(self)\n        TimeFrequencyScatteringBase.create_filters(self)\n\n\nTimeFrequencyScatteringJax._document()\n\n__all__ = [\"ScatteringJax1D\", \"TimeFrequencyScatteringJax\"]\n", "path": "kymatio/scattering1d/frontend/jax_frontend.py"}]}