Repository: kymatio/kymatio

Issue:
Scattering 2D doesn't work when using 2^J == image size (TensorFlow backend)
Similar issues have been reported before (#284, #363) and also fixed (#412) for Torch backend. However, still not working in the tensorflow backend:

test:
```python 
import numpy as np
from kymatio.tensorflow import Scattering2D

scattering = Scattering2D(J=5, shape=(32, 32))
test_im = np.ones((1,1,32,32))
test = scattering.scattering(test_im)
```

Gives error:
```
      7 scattering = Scattering2D(J=5, shape=(32, 32))
      8 test_im = np.ones((1,1,32,32))
----> 9 test = scattering.scattering(test_im)

~/.local/lib/python3.9/site-packages/kymatio/scattering2d/frontend/tensorflow_frontend.py in scattering(self, input)
     48             input = tf.reshape(input, tf.concat(((-1,), signal_shape), 0))
     49 
---> 50             S = scattering2d(input, self.pad, self.unpad, self.backend, self.J, self.L, self.phi, self.psi,
     51                              self.max_order, self.out_type)
     52 

~/.local/lib/python3.9/site-packages/kymatio/scattering2d/core/scattering2d.py in scattering2d(x, pad, unpad, backend, J, L, phi, psi, max_order, out_type)
     13     out_S_0, out_S_1, out_S_2 = [], [], []
     14 
---> 15     U_r = pad(x)
     16 
     17     U_0_c = fft(U_r, 'C2C')

~/.local/lib/python3.9/site-packages/kymatio/scattering2d/backend/tensorflow_backend.py in __call__(self, x)
     27         paddings = [[0, 0]] * len(x.shape[:-2])
     28         paddings += [[self.pad_size[0], self.pad_size[1]], [self.pad_size[2], self.pad_size[3]]]
---> 29         return tf.cast(tf.pad(x, paddings, mode="REFLECT"), tf.complex64)
     30 
     31 def unpad(in_):

~/.local/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
    151     except Exception as e:
    152       filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153       raise e.with_traceback(filtered_tb) from None
    154     finally:
    155       del filtered_tb

~/.local/lib/python3.9/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     56   try:
     57     ctx.ensure_initialized()
---> 58     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     59                                         inputs, attrs, num_outputs)
     60   except core._NotOkStatusException as e:

InvalidArgumentError: paddings must be less than the dimension size: 32, 32 not less than 32 [Op:MirrorPad]
```

(speculation) So possibly problems with the order of the padding being different in tensorflow from torch. 

Should also probably include some tests for these types of problems like the tests for implemented for Torch in #346 


Discussion:
This will be a bit involved on how to solve. according to [tensorflow documentation on MirrorPad](https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/mirror-pad); 

> Both `paddings[D, 0]` and `paddings[D, 1]` must be no greater than `input.dim_size(D)` (or `input.dim_size(D) - 1`) if `copy_border` is true (if false, respectively).

Sounds like we would need to do a symmetric padding to fix the bug, but I dont believe that is what we are supposed to do in terms of scattering... @janden @eickenberg @edouardoyallon how do we wanna reconcile this? 


[Google colab notebook replicating the issue](https://colab.research.google.com/drive/1DebjfI8gTAPFlBgTstlDDxXTwKdbrXoq?usp=sharing)

Edit: actually, it should be very straightforward to solve this. We simply replicate the [solution here](https://github.com/kymatio/kymatio/blob/dev/kymatio/scattering2d/backend/torch_backend.py#L48). 

Resolving patch:
diff --git a/kymatio/scattering2d/backend/tensorflow_backend.py b/kymatio/scattering2d/backend/tensorflow_backend.py
--- a/kymatio/scattering2d/backend/tensorflow_backend.py
+++ b/kymatio/scattering2d/backend/tensorflow_backend.py
@@ -16,11 +16,34 @@ def __init__(self, pad_size, input_size):
                 size of the original signal
         """
         self.pad_size = pad_size
+        self.input_size = input_size
 
     def __call__(self, x):
+        pad_size = list(self.pad_size)
+
+        # Clone to avoid passing on modifications.
+        new_pad_size = list(pad_size)
+
+        # This handles the case where the padding is equal to the image size.
+        if pad_size[0] == self.input_size[0]:
+            new_pad_size[0] -= 1
+            new_pad_size[1] -= 1
+        if pad_size[2] == self.input_size[1]:
+            new_pad_size[2] -= 1
+            new_pad_size[3] -= 1
+
         paddings = [[0, 0]] * len(x.shape[:-2])
-        paddings += [[self.pad_size[0], self.pad_size[1]], [self.pad_size[2], self.pad_size[3]]]
-        return tf.pad(x, paddings, mode="REFLECT")
+        paddings += [[new_pad_size[0], new_pad_size[1]], [new_pad_size[2], new_pad_size[3]]]
+
+        x_padded = tf.pad(x, paddings, mode="REFLECT")
+
+        # Again, special handling for when padding is the same as image size.
+        if pad_size[0] == self.input_size[0]:
+            x_padded = tf.concat([tf.expand_dims(x_padded[..., 1, :], axis=-2), x_padded, tf.expand_dims(x_padded[..., x_padded.shape[-2] -2, :], axis=-2)], axis=-2)
+        if pad_size[2] == self.input_size[1]:
+            x_padded = tf.concat([tf.expand_dims(x_padded[..., :, 1], axis=-1), x_padded, tf.expand_dims(x_padded[..., :,  x_padded.shape[-1]-2], axis=-1)], axis=-1)
+
+        return x_padded
 
 
 class TensorFlowBackend2D(TensorFlowBackend):


Test patch:
diff --git a/tests/scattering2d/test_tensorflow_scattering2d.py b/tests/scattering2d/test_tensorflow_scattering2d.py
--- a/tests/scattering2d/test_tensorflow_scattering2d.py
+++ b/tests/scattering2d/test_tensorflow_scattering2d.py
@@ -113,3 +113,10 @@ def test_inputs(self):
         with pytest.raises(RuntimeError) as ve:
             scattering = Scattering2D(10, shape=(10, 10), frontend='tensorflow')
         assert 'smallest dimension' in ve.value.args[0]
+
+class TestScatteringPad2DTensorFlow:
+    @pytest.mark.parametrize('backend', backends)
+    def test_Scattering2D_pad(self, backend):
+        scattering = Scattering2D(J=5, shape=(32, 32), pre_pad=False, frontend="tensorflow", backend=backend)
+        test_im = np.ones((1,1,32,32))
+        test = scattering(test_im)