Repository: google/flax

Issue:
flax.linen.Conv needs better error checking of 'padding' argument.
Hi!

The following code leads to mysterious error message `RuntimeError: UNKNOWN: -:4:130: error: expected '['` :

```
x = np.random.normal(size=(7, 48, 48, 96)).astype(np.float32)
model_def = nn.Conv(
    features=96, kernel_size=(7, 7),
    strides=(4, 4),
    padding=(3, 3))
model_state, conv_params = model_def.init({'params': jax.random.PRNGKey(42)}, x).pop('params')

out = model_def.apply({"params": conv_params}, x)
```

The mistake here is that I was using `padding=(3, 3)` instead of `padding=((3, 3), (3, 3))`, but the error message is not informative. It would be great if that could be improved. Ideally, a simpler padding spec. like `padding=(3, 3)` or even `padding=3` could directly be supported.


Discussion:

Resolving patch:
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -205,6 +205,33 @@ def _conv_dimension_numbers(input_shape):
   return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)
 
 
+PaddingLike = Union[str, int, Sequence[Union[int, Tuple[int, int]]]] 
+LaxPadding = Union[str, Sequence[Tuple[int, int]]]
+
+
+def canonicalize_padding(padding: PaddingLike, rank: int) -> LaxPadding:
+  """"Canonicalizes conv padding to a jax.lax supported format."""
+  if isinstance(padding, str):
+    return padding
+  if isinstance(padding, int):
+    return [(padding, padding)] * rank
+  if isinstance(padding, Sequence) and len(padding) == rank:
+    new_pad = []
+    for p in padding:
+      if isinstance(p, int):
+        new_pad.append((p, p))
+      elif isinstance(p, tuple) and len(p) == 2:
+        new_pad.append(p)
+      else:
+        break
+    if len(new_pad) == rank:
+      return new_pad
+  raise ValueError(
+    f'Invalid padding format: {padding}, should be str, int,'
+    f' or a sequence of len {rank} where each element is an'
+    f' int or pair of ints.')
+
+
 class _Conv(Module):
   """Convolution Module wrapping `lax.conv_general_dilated[_local]`.
 
@@ -218,7 +245,9 @@ class _Conv(Module):
     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.
+      spatial dimension. A single int is interpeted as applying the same padding
+      in all dims and passign a single int in a sequence causes the same padding
+      to be used on both sides.
     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
@@ -240,7 +269,7 @@ class _Conv(Module):
   features: int
   kernel_size: Sequence[int]
   strides: Union[None, int, Sequence[int]] = 1
-  padding: Union[str, Sequence[Tuple[int, int]]] = 'SAME'
+  padding: PaddingLike = 'SAME'
   input_dilation: Union[None, int, Sequence[int]] = 1
   kernel_dilation: Union[None, int, Sequence[int]] = 1
   feature_group_count: int = 1
@@ -307,8 +336,8 @@ def maybe_broadcast(x: Optional[Union[int, Sequence[int]]]) -> (
     input_dilation = maybe_broadcast(self.input_dilation)
     kernel_dilation = maybe_broadcast(self.kernel_dilation)
 
-    padding_lax: Union[str, Sequence[Tuple[int, int]]]
-    if self.padding == 'CIRCULAR':
+    padding_lax = canonicalize_padding(self.padding, len(kernel_size))
+    if padding_lax == 'CIRCULAR':
       kernel_size_dilated = [
           (k - 1) * d + 1 for k, d in zip(kernel_size, kernel_dilation)
       ]
@@ -317,8 +346,6 @@ def maybe_broadcast(x: Optional[Union[int, Sequence[int]]]) -> (
               [(0, 0)])
       inputs = jnp.pad(inputs, pads, mode='wrap')
       padding_lax = 'VALID'
-    else:
-      padding_lax = self.padding
 
     dimension_numbers = _conv_dimension_numbers(inputs.shape)
     in_features = inputs.shape[-1]
@@ -429,7 +456,9 @@ class ConvTranspose(Module):
     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.
+      spatial dimension. A single int is interpeted as applying the same padding
+      in all dims and passign a single int in a sequence causes the same padding
+      to be used on both sides.
     kernel_dilation: `None`, or a sequence of `n` integers, giving the
       dilation factor to apply in each spatial dimension of the convolution
       kernel. Convolution with kernel dilation is also known as 'atrous
@@ -445,7 +474,7 @@ class ConvTranspose(Module):
   features: int
   kernel_size: Union[int, Tuple[int, ...]]
   strides: Optional[Tuple[int, ...]] = None
-  padding: Union[str, Sequence[Tuple[int, int]]] = 'SAME'
+  padding: PaddingLike = 'SAME'
   kernel_dilation: Optional[Sequence[int]] = None
   use_bias: bool = True
   dtype: Dtype = jnp.float32
@@ -492,11 +521,9 @@ def __call__(self, inputs: Array) -> Array:
                         self.param_dtype)
     kernel = jnp.asarray(kernel, self.dtype)
 
-    padding_lax: Union[str, Sequence[Tuple[int, int]]]
-    if self.padding == 'CIRCULAR':
+    padding_lax = canonicalize_padding(self.padding, len(kernel_size))
+    if padding_lax == 'CIRCULAR':
       padding_lax = 'VALID'
-    else:
-      padding_lax = self.padding
 
     y = lax.conv_transpose(
         inputs,


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
@@ -15,6 +15,7 @@
 """Tests for flax.deprecated.nn.linear."""
 
 import functools
+from multiprocessing.sharedctypes import Value
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -842,6 +843,19 @@ def __call__(self, x):
         }})
     self.assertEqual(y.shape, (8, 6))
 
+  def test_canonicalize_padding(self):
+    def test_pad(pad, rank, expected=None):
+      if expected is None:
+        with self.assertRaises(ValueError):
+          nn.linear.canonicalize_padding(pad, rank)
+      else:
+        self.assertEqual(nn.linear.canonicalize_padding(pad, rank), expected)
+    test_pad("SAME", 2, "SAME")
+    test_pad(2, 3, [(2, 2), (2, 2), (2, 2)])
+    test_pad((2, 2), 3)
+    test_pad((2, 2), 1)
+    test_pad([1, (2, 3)], 2, [(1, 1), (2, 3)])
+    test_pad([None, (1, 2)], 2)
 
 if __name__ == '__main__':
   absltest.main()