Repository: google/flax

Issue:
Circular padding in convolutional neural networks
### Description of the model to be implemented

In many areas such as physics, it is convenient to have convolutional layers with periodic boundary conditions (e.g. see [netket](https://github.com/netket/netket))

Therefore, it would be nice to add a "CIRCULAR" padding option to convolutional layers, just as they do in [neural-tangents](https://neural-tangents.readthedocs.io/en/latest/_modules/neural_tangents/stax.html#Conv).

### Dataset the model could be trained on

1D or 2D data. Maybe MNIST images.

### Specific points to consider

None in particular. Just as an example, suppose that one has the 1D data [1,2,3,4,5] and one has filters of size 3, and a stride of 3. The idea is then that two filter operations are carried out. The first one will use [1,2,3], and the second one will use [4,5,1].

### Reference implementations in other frameworks

neural-tangents has replaced stax's GeneralConv by a Conv layer, which has this padding option, and further does not require to provide directly the XLA's `dimension_numbers`.


Discussion:
I think it would be quite nice to add this, since it doesn't seem to complicate the API much (no additional parameters etc). @levskaya what do you think of this proposal? I recall you were involved in a discussion around this before, and I'm curious whether you think it makes sense to add this.
It would be even nicer if the jax conv op would support this out of the box. They already have 'same' and 'valid'.
If this is still relevant, I'd be happy to raise a PR, reusing the code from https://github.com/google/flax/issues/903#issue-789095219 and adding some tests
I'd love it if you do that Grisha!

On Wed, 3 Nov 2021 at 17:55, Grisha Oryol ***@***.***> wrote:

> If this is still relevant, I'd be happy to raise a PR, reusing the code
> from #903 (comment)
> <https://github.com/google/flax/issues/903#issue-789095219> and adding
> some tests
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/google/flax/issues/971#issuecomment-960260494>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AHQJA72FDQBHK5DPXFZCKLTUKG4XLANCNFSM4W4YKQGQ>
> .
> Triage notifications on the go with GitHub Mobile for iOS
> <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
> or Android
> <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.
>
>


-- 
Best wishes,

Vladimir Vargas-Calderón
<https://www.researchgate.net/profile/Vladimir_Vargas-Calderon>
PhD Physics Student @ Universidad Nacional de Colombia

-- 
*Aviso legal:* El contenido de este mensaje y los archivos adjuntos son 
confidenciales y de uso exclusivo de la Universidad Nacional de Colombia. 
Se encuentran dirigidos sólo para el uso del destinatario al cual van 
enviados. La reproducción, lectura y/o copia se encuentran prohibidas a 
cualquier persona diferente a este y puede ser ilegal. Si usted lo ha 
recibido por error, infórmenos y elimínelo de su correo. Los Datos 
Personales serán tratados conforme a la Ley 1581 de 2012 y a nuestra 
Política de Datos Personales que podrá consultar en la página web 
www.unal.edu.co <http://www.unal.edu.co/>.* *Las opiniones, informaciones, 
conclusiones y cualquier otro tipo de dato contenido en este correo 
electrónico, no relacionados con la actividad de la Universidad Nacional de 
Colombia, se entenderá como personales y de ninguna manera son avaladas por 
la Universidad.


Resolving patch:
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -201,9 +201,9 @@ class Conv(Module):
       be a sequence of integers.
     strides: an integer or a sequence of `n` integers, representing the
       inter-window strides (default: 1).
-    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
-      of `n` `(low, high)` integer pairs that give the padding to apply before
-      and after each spatial dimension.
+    padding: either the string `'SAME'`, the string `'VALID'`, the string 'CIRCULAR'` (periodic boundary conditions),
+      or a sequence of `n` `(low, high)` integer pairs that give the padding to apply
+      before and after each spatial dimension.
     input_dilation: an integer or a sequence of `n` integers, giving the
       dilation factor to apply in each spatial dimension of `inputs` (default: 1).
       Convolution with input dilation `d` is equivalent to transposed
@@ -282,12 +282,20 @@ def maybe_broadcast(x):
     kernel = self.param('kernel', self.kernel_init, kernel_shape)
     kernel = jnp.asarray(kernel, self.dtype)
 
+    if self.padding == 'CIRCULAR':
+      kernel_size_dilated = [(k - 1) * d + 1 for k, d in zip(kernel_size, kernel_dilation)]
+      pads = [(0, 0)] + [((k - 1) // 2, k // 2) for k in kernel_size_dilated] + [(0, 0)]
+      inputs = jnp.pad(inputs, pads, mode='wrap')
+      padding_lax = 'VALID'
+    else:
+      padding_lax = self.padding
+
     dimension_numbers = _conv_dimension_numbers(inputs.shape)
     y = lax.conv_general_dilated(
         inputs,
         kernel,
         strides,
-        self.padding,
+        padding_lax,
         lhs_dilation=input_dilation,
         rhs_dilation=kernel_dilation,
         dimension_numbers=dimension_numbers,
@@ -313,8 +321,8 @@ class ConvTranspose(Module):
       be a sequence of integers.
     strides: a sequence of `n` integers, representing the inter-window
       strides.
-    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
-      of `n` `(low, high)` integer pairs that give the padding to apply before
+    padding: either the string `'SAME'`, the string `'VALID'`, the string 'CIRCULAR'` (periodic boundary conditions),
+      or a sequence of `n` `(low, high)` integer pairs that give the padding to apply before
       and after each spatial dimension.
     kernel_dilation: `None`, or a sequence of `n` integers, giving the
       dilation factor to apply in each spatial dimension of the convolution
@@ -372,13 +380,49 @@ def __call__(self, inputs: Array) -> Array:
     kernel = self.param('kernel', self.kernel_init, kernel_shape)
     kernel = jnp.asarray(kernel, self.dtype)
 
+    if self.padding == 'CIRCULAR':
+      padding_lax = 'VALID'
+    else:
+      padding_lax = self.padding
+
     y = lax.conv_transpose(inputs,
                            kernel,
                            strides,
-                           self.padding,
+                           padding_lax,
                            rhs_dilation=self.kernel_dilation,
                            precision=self.precision)
 
+    if self.padding == "CIRCULAR":
+      # For circular padding, we need to identify the size of the final output
+      # ("period") along each spatial dimension, pad each dimension to an
+      # integer number of periods, and wrap the array periodically around each
+      # dimension. Padding should be done in such a way that the start of the
+      # original input data inside the padded array is located at integer
+      # number of periods - otherwise the result would be circularly shifted.
+      
+      # Compute period along each spatial dimension - it's input size scaled
+      # by the stride.
+      scaled_x_dims = [
+        x_dim * stride for x_dim, stride in zip(inputs.shape[1:-1], strides)
+      ]
+      # Compute difference between the current size of y and the final output
+      # size, and complement this difference to 2 * period - that gives how
+      # much we need to pad.
+      size_diffs = [
+        -(y_dim - x_dim) % (2 * x_dim)
+        for y_dim, x_dim in zip(y.shape[1:-1], scaled_x_dims)
+      ]
+      # Divide the padding equaly between left and right. The choice to put
+      # "+1" on the left (and not on the right) represents a convention for
+      # aligning even-sized kernels.
+      total_pad = [((size_diff + 1) // 2, size_diff // 2) for size_diff in size_diffs]
+      y = np.pad(y, [(0, 0)] + total_pad + [(0, 0)])
+      # Wrap the result periodically around each spatial dimension,
+      # one by one.
+      for i in range(1, y.ndim - 1):
+        y = y.reshape(y.shape[:i] + (-1, scaled_x_dims[i - 1]) + y.shape[i + 1:])
+        y = y.sum(axis=i)
+
     if is_single_input:
       y = jnp.squeeze(y, axis=0)
     if self.use_bias:


Test patch:
diff --git a/tests/linen/linen_linear_test.py b/tests/linen/linen_linear_test.py
--- a/tests/linen/linen_linear_test.py
+++ b/tests/linen/linen_linear_test.py
@@ -204,6 +204,174 @@ def test_group_conv(self):
     self.assertEqual(initial_params['params']['kernel'].shape, (3, 2, 4))
     np.testing.assert_allclose(y, np.full((1, 6, 4), 7.))
 
+  @parameterized.product(
+      n_batch=(1, 3),
+      n_features=(1, 2),
+      kernel_size=(1, 2, 3, 9),
+      n_input_features=(1, 3),
+      input_size=(1, 8, 16),
+  )
+  def test_circular_conv_1d_constant(
+          self, n_batch, n_features, kernel_size, n_input_features, input_size
+  ):
+      """
+      Test 1D convolution with circular padding: filter with all elements equal to 1
+      applied on an input with all elements equal to 1.
+      Result should have the same shape as input (except for the feature dimension) and
+      have all elements equal to n_input_features * kernel_lin_size
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = jnp.ones((n_batch, input_size, n_input_features))
+      conv_module = nn.Conv(
+          features=n_features,
+          kernel_size=(kernel_size,),
+          padding='CIRCULAR',
+          kernel_init=initializers.ones,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(
+          initial_params['params']['kernel'].shape,
+          (kernel_size, n_input_features, n_features),
+      )
+      correct_ans = np.full(
+          (n_batch, input_size, n_features), kernel_size * n_input_features
+      )
+      np.testing.assert_allclose(y, correct_ans)
+
+  @parameterized.product(
+      n_batch=(1, 3),
+      n_features=(1, 2, 10),
+      kernel_lin_size=(1, 2, 3, 9),
+      n_input_features=(1, 5),
+      input_x_size=(14,),
+      input_y_size=(5, 10),
+  )
+  def test_circular_conv_2d_constant(
+          self,
+          n_batch,
+          n_features,
+          kernel_lin_size,
+          n_input_features,
+          input_x_size,
+          input_y_size,
+  ):
+      """
+      Test 2D convolution with circular padding: square filter with all elements equal to 1
+      applied on an input with all elements equal to 1.
+      Result should have the same shape as input (except for the feature dimension) and
+      have all elements equal to n_input_features * kernel_lin_size^2
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = jnp.ones((n_batch, input_x_size, input_y_size, n_input_features))
+      conv_module = nn.Conv(
+          features=n_features,
+          kernel_size=(kernel_lin_size, kernel_lin_size),
+          padding='CIRCULAR',
+          kernel_init=initializers.ones,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(
+          initial_params['params']['kernel'].shape,
+          (kernel_lin_size, kernel_lin_size, n_input_features, n_features),
+      )
+      correct_ans = np.full(
+          (n_batch, input_x_size, input_y_size, n_features),
+          kernel_lin_size * kernel_lin_size * n_input_features,
+      )
+      np.testing.assert_allclose(y, correct_ans)
+
+  def test_circular_conv_1d_custom(self):
+      """
+      Test 1d convolution with circular padding and a stride
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = np.arange(1, 6)
+      x = np.expand_dims(x, (0, 2))
+      kernel = np.array((1, 2, 1))
+      kernel = np.expand_dims(kernel, (1, 2))
+
+      conv_module = nn.Conv(
+          features=1,
+          kernel_size=(3,),
+          strides=(3,),
+          padding='CIRCULAR',
+          kernel_init=lambda *_: kernel,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(initial_params['params']['kernel'].shape, (3, 1, 1))
+      # Compare with manually computed convolution
+      correct_ans = np.array((5 + 2 * 1 + 2, 3 + 2 * 4 + 5))
+      correct_ans = np.expand_dims(correct_ans, (0, 2))
+      np.testing.assert_allclose(y, correct_ans)
+
+
+  def test_circular_conv_1d_dilation(self):
+      """
+      Test 1d convolution with circular padding and kernel dilation
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = np.arange(1, 6)
+      x = np.expand_dims(x, (0, 2))
+      kernel = np.array((1, 2, 1))
+      kernel = np.expand_dims(kernel, (1, 2))
+
+      conv_module = nn.Conv(
+          features=1,
+          kernel_size=(3,),
+          padding='CIRCULAR',
+          kernel_init=lambda *_: kernel,
+          bias_init=initializers.zeros,
+          kernel_dilation=(3,)
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(initial_params['params']['kernel'].shape, (3, 1, 1))
+      # Compare with manually computed convolution
+      correct_ans = np.array((3 + 2 * 1 + 4, 4 + 2 * 2 + 5, 5 + 2 * 3 + 1, 1 + 2 * 4 + 2, 2 + 2 * 5 + 3))
+      correct_ans = np.expand_dims(correct_ans, (0, 2))
+      np.testing.assert_allclose(y, correct_ans)
+
+  def test_circular_conv_2d_custom(self):
+      """
+      Test 2d convolution with circular padding on a 3x3 example
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = np.array(((1, 2, 3),
+                    (4, 5, 6),
+                    (7, 8, 9)))
+      x = np.expand_dims(x, (0, 3))
+      kernel = np.array(((0, 1, 0),
+                         (1, 2, 1),
+                         (0, 1, 0)))
+      kernel = np.expand_dims(kernel, (2, 3))
+
+      conv_module = nn.Conv(
+          features=1,
+          kernel_size=(3, 3),
+          padding='CIRCULAR',
+          kernel_init=lambda *_: kernel,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 1, 1))
+      # Compare with manually computed convolution
+      correct_ans = np.array(
+          (
+              (2 * 1 + 7 + 2 + 4 + 3, 2 * 2 + 8 + 3 + 5 + 1, 2 * 3 + 9 + 1 + 6 + 2),
+              (2 * 4 + 1 + 5 + 7 + 6, 2 * 5 + 2 + 6 + 8 + 4, 2 * 6 + 3 + 4 + 9 + 5),
+              (2 * 7 + 4 + 8 + 1 + 9, 2 * 8 + 5 + 9 + 2 + 7, 2 * 9 + 6 + 7 + 3 + 8),
+          )
+      )
+      correct_ans = np.expand_dims(correct_ans, (0, 3))
+      np.testing.assert_allclose(y, correct_ans)
+
   def test_conv_transpose(self):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((1, 8, 3))
@@ -252,6 +420,202 @@ def test_single_input_conv_transpose(self):
                               [ 4.,  4.,  4.,  4.]])
     np.testing.assert_allclose(y, correct_ans)
 
+  @parameterized.product(
+      n_batch=(1, 3),
+      n_features=(1, 2),
+      kernel_size=(1, 2, 3, 9),
+      n_input_features=(1, 3),
+      input_size=(1, 8, 16),
+  )
+  def test_circular_conv_transpose_1d_constant(
+          self, n_batch, n_features, kernel_size, n_input_features, input_size
+  ):
+      """
+      Test 1D transposed convolution with circular padding: filter with all elements equal to 1
+      applied on an input with all elements equal to 1.
+      Result should have the same shape as input (except for the feature dimension) and
+      have all elements equal to n_input_features * kernel_lin_size
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = jnp.ones((n_batch, input_size, n_input_features))
+      conv_module = nn.ConvTranspose(
+          features=n_features,
+          kernel_size=(kernel_size,),
+          padding="CIRCULAR",
+          kernel_init=initializers.ones,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(
+          initial_params["params"]["kernel"].shape,
+          (kernel_size, n_input_features, n_features),
+      )
+      correct_ans = np.full(
+          (n_batch, input_size, n_features), kernel_size * n_input_features
+      )
+      np.testing.assert_allclose(y, correct_ans)
+
+  @parameterized.product(
+      n_batch=(1, 3),
+      n_features=(1, 2, 10),
+      kernel_lin_size=(1, 2, 3, 9),
+      n_input_features=(1, 5),
+      input_x_size=(14,),
+      input_y_size=(5, 10),
+  )
+  def test_circular_conv_transpose_2d_constant(
+          self,
+          n_batch,
+          n_features,
+          kernel_lin_size,
+          n_input_features,
+          input_x_size,
+          input_y_size,
+  ):
+      """
+      Test 2D transposed convolution with circular padding: square filter with all elements equal to 1
+      applied on an input with all elements equal to 1.
+      Result should have the same shape as input (except for the feature dimension) and
+      have all elements equal to n_input_features * kernel_lin_size^2
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = jnp.ones((n_batch, input_x_size, input_y_size, n_input_features))
+      conv_module = nn.ConvTranspose(
+          features=n_features,
+          kernel_size=(kernel_lin_size, kernel_lin_size),
+          padding="CIRCULAR",
+          kernel_init=initializers.ones,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(
+          initial_params["params"]["kernel"].shape,
+          (kernel_lin_size, kernel_lin_size, n_input_features, n_features),
+      )
+      correct_ans = np.full(
+          (n_batch, input_x_size, input_y_size, n_features),
+          kernel_lin_size * kernel_lin_size * n_input_features,
+      )
+      np.testing.assert_allclose(y, correct_ans)
+
+  def test_circular_conv_transpose_1d_custom(self):
+      """
+      Test 1d transposed convolution with circular padding and a stride
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = np.arange(1, 6)
+      x = np.expand_dims(x, (0, 2))
+      kernel = np.array((1, 2, 1))
+      kernel = np.expand_dims(kernel, (1, 2))
+
+      conv_module = nn.ConvTranspose(
+          features=1,
+          kernel_size=(3,),
+          strides=(3,),
+          padding="CIRCULAR",
+          kernel_init=lambda *_: kernel,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(initial_params["params"]["kernel"].shape, (3, 1, 1))
+      # Compare with manually computed convolution
+      correct_ans = np.array(
+          (1 * 1, 1 * 2, 1 * 1,
+           2 * 1, 2 * 2, 2 * 1,
+           3 * 1, 3 * 2, 3 * 1,
+           4 * 1, 4 * 2, 4 * 1,
+           5 * 1, 5 * 2, 5 * 1,
+           )
+      )
+      correct_ans = np.expand_dims(correct_ans, (0, 2))
+      np.testing.assert_allclose(y, correct_ans)
+
+  def test_circular_conv_transpose_2d_custom(self):
+      """
+      Test 2d transposed convolution with circular padding on a 3x3 example
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = np.array(
+          (
+              (1, 2, 3),
+              (4, 5, 6),
+              (7, 8, 9),
+          )
+      )
+      x = np.expand_dims(x, (0, 3))
+      kernel = np.array(
+          (
+              (0, 1, 0),
+              (1, 2, 1),
+              (0, 1, 0)
+          )
+      )
+      kernel = np.expand_dims(kernel, (2, 3))
+
+      conv_module = nn.ConvTranspose(
+          features=1,
+          kernel_size=(3, 3),
+          padding="CIRCULAR",
+          kernel_init=lambda *_: kernel,
+          bias_init=initializers.zeros,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(initial_params["params"]["kernel"].shape, (3, 3, 1, 1))
+      # Compare with manually computed convolution
+      correct_ans = np.array(
+          (
+              (18, 21, 24),
+              (27, 30, 33),
+              (36, 39, 42),
+          )
+      )
+      correct_ans = np.expand_dims(correct_ans, (0, 3))
+      np.testing.assert_allclose(y, correct_ans)
+
+  def test_circular_conv_transpose_2d_custom_bias(self):
+      """
+      Test 2d transposed convolution with circular padding on a 2x2 example with bias
+      """
+      rng = dict(params=random.PRNGKey(0))
+      x = np.array(
+          (
+              (1, 2),
+              (3, 4)
+          )
+      )
+      x = np.expand_dims(x, (0, 3))
+      kernel = np.array(
+          (
+              (1, 2),
+              (3, 4),
+          )
+      )
+      kernel = np.expand_dims(kernel, (2, 3))
+
+      conv_module = nn.ConvTranspose(
+          features=1,
+          kernel_size=(2, 2),
+          padding="CIRCULAR",
+          kernel_init=lambda *_: kernel,
+          bias_init=initializers.ones,
+      )
+      y, initial_params = conv_module.init_with_output(rng, x)
+
+      self.assertEqual(initial_params["params"]["kernel"].shape, (2, 2, 1, 1))
+      # Compare with manually computed convolution
+      correct_ans = np.array(
+          (
+              (21, 23),
+              (29, 31),
+          )
+      )
+      correct_ans = np.expand_dims(correct_ans, (0, 3))
+      np.testing.assert_allclose(y, correct_ans)
+
   def test_int_kernel_size(self):
     conv = nn.Conv(features=4, kernel_size=3)
     x = jnp.ones((8, 3))