Repository: kymatio/kymatio

Issue:
Make inputs dimension-agnostic (1D/2D)
In other words, make `Scattering*D` operators act only on last dimension without forcing a specific convention on previous dimensions. The output would have the same shape, but this the time/spatial dimension subsampled and another (scattering) dimension added right before it. So we'd have

- in 1D, `batch_shape + (T,)` input gives `batch_shape + (C, T/2**J)` output, and
- in 2D, `batch_shape + (N, M)` input gives `batch_shape + (C, N/2**J, M/2**J)` output,

where `batch_shape` is some (possibly empty) tuple specifying the batch dimensions that are not to be processed and `C` is the number of scattering coefficients.

The 3D case is more complicated since it follows a different output schema compared to the 1D and 2D, which will be addressed in a separate issue. That being said, it would be good if it could mimic the above behavior.


Discussion:
Thanks for raising this Joakim. For future reference, this originated in PR #215 which closed issue #213, which itself was derived off #158.
Quoting Semver 7:
> Minor version Y (x.Y.z | x > 0) MUST be incremented if new, backwards compatible functionality is introduced to the public API. 

Therefore this change should be marked as `0.2.0`, not `0.1.0-beta`. The `beta` release will only add bugfixes and improvements to the docs.
Even if it's backwards-compatible, huh? Didn't know that.

Ok so unless we can push this into `0.1.0-alpha`, `0.2.0` it is. It honestly shouldn't be that hard.
https://semver.org is an insightful read.
during `0.y.z` phase:
* backwards-incompatible API changes push to `0.(y+1).z`
* backwards-compatible API changes push to `0.(y+1).z`
* bugfix changes push to `0.y.(z+1)`

during `x.y.z` phase where `x>0`:
* backwards-incompatible API changes push to `(x+1).y.z`
* backwards-compatible API changes push to `x.(y+1).z`
* bugfix changes push to `x.y.(z+1)`
Ok makes sense. Good thing we don't have `x>0` then. Sounds like a pain....

Resolving patch:
diff --git a/kymatio/scattering2d/scattering2d.py b/kymatio/scattering2d/scattering2d.py
--- a/kymatio/scattering2d/scattering2d.py
+++ b/kymatio/scattering2d/scattering2d.py
@@ -146,6 +146,10 @@ def forward(self, input):
         if not torch.is_tensor(input):
             raise(TypeError('The input should be a torch.cuda.FloatTensor, a torch.FloatTensor or a torch.DoubleTensor'))
 
+        if len(input.shape) < 2:
+            raise (RuntimeError('Input tensor must have at least two '
+                'dimensions'))
+
         if (not input.is_contiguous()):
             raise (RuntimeError('Tensor must be contiguous!'))
 
@@ -155,8 +159,10 @@ def forward(self, input):
         if ((input.size(-1) != self.N_padded or input.size(-2) != self.M_padded) and self.pre_pad):
             raise (RuntimeError('Padded tensor must be of spatial size (%i,%i)!' % (self.M_padded, self.N_padded)))
 
-        if (input.dim() != 4):
-            raise (RuntimeError('Input tensor must be 4D'))
+        batch_shape = input.shape[:-2]
+        signal_shape = input.shape[-2:]
+
+        input = input.reshape((-1, 1) + signal_shape)
 
         J = self.J
         phi = self.Phi
@@ -215,6 +221,9 @@ def forward(self, input):
                         S[..., n, :, :] = unpad(U_J_r)
                         n = n + 1
 
+        scattering_shape = S.shape[-3:]
+        S = S.reshape(batch_shape + scattering_shape)
+
         return S
 
     def __call__(self, input):


Test patch:
diff --git a/kymatio/scattering2d/tests/test_scattering2d.py b/kymatio/scattering2d/tests/test_scattering2d.py
--- a/kymatio/scattering2d/tests/test_scattering2d.py
+++ b/kymatio/scattering2d/tests/test_scattering2d.py
@@ -2,6 +2,7 @@
 
 import os
 import torch
+import pytest
 from kymatio.scattering2d import Scattering2D
 from kymatio.scattering2d import backend
 
@@ -184,3 +185,41 @@ def test_Scattering2D():
                 Sg = scattering(x)
             assert (Sg - S).abs().max() < 1e-6
 
+
+def test_batch_shape_agnostic():
+    J = 3
+    L = 8
+    shape = (32, 32)
+
+    shape_ds = tuple(n // 2 **J for n in shape)
+
+    S = Scattering2D(J, shape, L)
+
+    with pytest.raises(RuntimeError) as ve:
+        S(torch.zeros(()))
+    assert "at least two" in ve.value.args[0]
+
+    with pytest.raises(RuntimeError) as ve:
+        S(torch.zeros((32, )))
+    assert "at least two" in ve.value.args[0]
+
+    x = torch.zeros(shape)
+    Sx = S(x)
+
+    assert len(Sx.shape) == 3
+    assert Sx.shape[-2:] == shape_ds
+
+    n_coeffs = Sx.shape[-3]
+
+    test_shapes = ((1,) + shape, (2,) + shape, (2, 2) + shape,
+                   (2, 2, 2) + shape)
+
+    for test_shape in test_shapes:
+        x = torch.zeros(test_shape)
+
+        Sx = S(x)
+
+        assert len(Sx.shape) == len(test_shape) + 1
+        assert Sx.shape[-2:] == shape_ds
+        assert Sx.shape[-3] == n_coeffs
+        assert Sx.shape[:-3] == test_shape[:-2]