Repository: kymatio/kymatio

Issue:
JTFS needs a JAX frontend
v0.4.0

Already addressed by #1001 (@cyrusvahidi)


Discussion:

Resolving patch:
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 @@ class ScatteringJax1D(ScatteringJax, ScatteringNumPy1D):
     # 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"]


Test patch:
diff --git a/tests/scattering1d/test_timefrequency_scattering.py b/tests/scattering1d/test_timefrequency_scattering.py
--- a/tests/scattering1d/test_timefrequency_scattering.py
+++ b/tests/scattering1d/test_timefrequency_scattering.py
@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 import torch, tensorflow as tf
+import jax.numpy as jnp
+from jax import device_put
 
 import kymatio
 from kymatio import TimeFrequencyScattering
@@ -130,14 +132,14 @@ def test_jtfs_create_filters(backend):
     )
 
 
-frontends = ["numpy", "torch", "tensorflow", "sklearn"]
+frontends = ["numpy", "torch", "tensorflow", "sklearn", "jax"]
 @pytest.mark.parametrize("frontend", frontends)
 def test_time_scattering_widthfirst(frontend):
     """Checks that width-first and depth-first algorithms have same output."""
     J = 5
     shape = (1024,)
     S = kymatio.Scattering1D(J, shape, frontend=frontend)
-    x = torch.zeros(shape)
+    x = torch.zeros(shape) if frontend == "torch" else np.zeros(shape)
     x[shape[0] // 2] = 1
     x_shape = S.backend.shape(x)
     _, signal_shape = x_shape[:-1], x_shape[-1:]
@@ -163,7 +165,7 @@ def test_time_scattering_widthfirst(frontend):
     S1_depth = S.backend.stack(
         [S1_depth[key] for key in sorted(S1_depth.keys()) if len(key) == 1]
     )
-    if frontend not in ["numpy", "sklearn"]:
+    if frontend not in ["numpy", "sklearn", "jax"]:
         S1_width = S1_width.numpy()
         S1_depth = S1_depth.numpy()
     assert np.allclose(S1_width, S1_depth)
@@ -177,7 +179,7 @@ def test_frequency_scattering(frontend):
     shape = (4096,)
     Q = 8
     S = kymatio.Scattering1D(J=J, Q=Q, shape=shape, frontend=frontend)
-    x = torch.zeros(shape)
+    x = np.zeros(shape)
     x[shape[0] // 2] = 1
     x_shape = S.backend.shape(x)
     _, signal_shape = x_shape[:-1], x_shape[-1:]
@@ -494,7 +496,7 @@ def test_differentiability_jtfs_tensorflow(random_state=42):
 def test_jtfs_numpy_and_sklearn(frontend):
     # Test __init__
     kwargs = {"J": 8, "J_fr": 3, "shape": (1024,), "Q": 3}
-    x = torch.zeros(kwargs["shape"])
+    x = np.zeros(kwargs["shape"])
     x[kwargs["shape"][0] // 2] = 1
 
     # Local averaging
@@ -581,6 +583,23 @@ def test_jtfs_torch_tf_frontends(frontend):
     assert Sx.ndim == 3
 
 
+frontends = ["jax"]
+@pytest.mark.parametrize("frontend", frontends)
+def test_jtfs_jax_frontend(frontend):
+    # Test __init__
+    kwargs = {"J": 8, "J_fr": 3, "shape": (8192,), "Q": 3}
+    x = np.zeros(kwargs["shape"])
+    x[kwargs["shape"][0] // 2] = 1
+    x = device_put(jnp.asarray(x))
+
+    # Local averaging
+    S = TimeFrequencyScattering(frontend=frontend, format="joint", **kwargs)
+    assert S.F == (2**S.J_fr)
+    Sx = S(x)
+    assert Sx.ndim == 3
+    assert S.backend.name == 'jax'
+
+
 frontends = ["torch", "tensorflow"]
 @pytest.mark.parametrize("frontend", frontends)
 def test_F(frontend):