Repository: google/flax

Issue:
Surprising behaviour for integer kernel_size in linen.Conv 
I was quite surprised with how `linen.Conv` treats `int` as an argument for `kernel_size`

```
key1, key2 = jax.random.split(jax.random.PRNGKey(0), 2)
image = jax.random.normal(key1, (8, 256, 256, 3))
conv = flax.linen.Conv(features=48, kernel_size=5)
params = conv.init(key2, image)
```

This errors on the last line with
```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
google3/third_party/py/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
     55     try:
---> 56         return getattr(obj, method)(*args, **kwds)
     57 

AttributeError: 'tuple' object has no attribute 'take'

During handling of the above exception, another exception occurred:

IndexError                                Traceback (most recent call last)
11 frames
google3/third_party/py/numpy/core/fromnumeric.py in _wrapit(obj, method, *args, **kwds)
     44     except AttributeError:
     45         wrap = None
---> 46     result = getattr(asarray(obj), method)(*args, **kwds)
     47     if wrap:
     48         if not isinstance(result, mu.ndarray):

IndexError: index 3 is out of bounds for size 3
```

And I can see it creates a kernel for only 1 spatial dimension.

So it seems it says `5 == (5,)`

Changing to `conv = flax.linen.Conv(features=48, kernel_size=(5,5))` fixes the error.

Overall I find:
a) the error is cryptic
b) if this is the desired behaviour, perhaps it should only accept kernel_size as a sequence of ints?



Discussion:
Your input is of shape `(8, 256, 256, 3)`, and when you specify a 1D kernel you are applying a 1D convolution. This doesn't work on your input shape, which expects 2D convolutions (your have 2 spatial dimensions). So you can fix this by reducing your input to one special dimension:

```python
from flax import linen as nn
import jax

key1, key2 = jax.random.split(jax.random.PRNGKey(0), 2)
image = jax.random.normal(key1, (8, 256, 4))
conv = nn.Conv(features=48, kernel_size=3)
params = conv.init(key2, image)
```

Closing this for now but please re-open if you think I missed something!
I think we want to keep this open as we want to get rid of flax's magic integer --> tuple conversion here

Resolving patch:
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -199,19 +199,19 @@ class Conv(Module):
     kernel_size: shape of the convolutional kernel. For 1D convolution,
       the kernel size can be passed as an integer. For all other cases, it must
       be a sequence of integers.
-    strides: a sequence of `n` integers, representing the inter-window
-      strides.
+    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.
-    input_dilation: `None`, or a sequence of `n` integers, giving the
-      dilation factor to apply in each spatial dimension of `inputs`.
+    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
       convolution with stride `d`.
-    kernel_dilation: `None`, or a sequence of `n` integers, giving the
+    kernel_dilation: an integer 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
-      convolution'.
+      kernel (default: 1). Convolution with kernel dilation
+      is also known as 'atrous convolution'.
     feature_group_count: integer, default 1. If specified divides the input
       features into groups.
     use_bias: whether to add a bias to the output (default: True).
@@ -222,11 +222,11 @@ class Conv(Module):
     bias_init: initializer for the bias.
   """
   features: int
-  kernel_size: Union[int, Iterable[int]]
-  strides: Optional[Iterable[int]] = None
+  kernel_size: Iterable[int]
+  strides: Union[None, int, Iterable[int]] = 1
   padding: Union[str, Iterable[Tuple[int, int]]] = 'SAME'
-  input_dilation: Optional[Iterable[int]] = None
-  kernel_dilation: Optional[Iterable[int]] = None
+  input_dilation: Union[None, int, Iterable[int]] = 1
+  kernel_dilation: Union[None, int, Iterable[int]] = 1
   feature_group_count: int = 1
   use_bias: bool = True
   dtype: Dtype = jnp.float32
@@ -248,16 +248,28 @@ def __call__(self, inputs: Array) -> Array:
     inputs = jnp.asarray(inputs, self.dtype)
 
     if isinstance(self.kernel_size, int):
-      kernel_size = (self.kernel_size,)
+      raise TypeError('The kernel size must be specified as a'
+                      ' tuple/list of integers (eg.: [3, 3]).')
     else:
-      kernel_size = self.kernel_size
+      kernel_size = tuple(self.kernel_size)
+
+    def maybe_broadcast(x):
+      if x is None:
+        # backward compatibility with using None as sentinel for
+        # broadcast 1
+        x = 1
+      if isinstance(x, int):
+        return (x,) * len(kernel_size)
+      return x
 
     is_single_input = False
     if inputs.ndim == len(kernel_size) + 1:
       is_single_input = True
       inputs = jnp.expand_dims(inputs, axis=0)
 
-    strides = self.strides or (1,) * (inputs.ndim - 2)
+    strides = maybe_broadcast(self.strides)  # self.strides or (1,) * (inputs.ndim - 2)
+    input_dilation = maybe_broadcast(self.input_dilation)
+    kernel_dilation = maybe_broadcast(self.kernel_dilation)
 
     in_features = inputs.shape[-1]
     assert in_features % self.feature_group_count == 0
@@ -272,8 +284,8 @@ def __call__(self, inputs: Array) -> Array:
         kernel,
         strides,
         self.padding,
-        lhs_dilation=self.input_dilation,
-        rhs_dilation=self.kernel_dilation,
+        lhs_dilation=input_dilation,
+        rhs_dilation=kernel_dilation,
         dimension_numbers=dimension_numbers,
         feature_group_count=self.feature_group_count,
         precision=self.precision)


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
@@ -161,13 +161,12 @@ def test_dense_general_vs_numpy(self, axis, batch_dims, einsum_expr):
     target = np.einsum(einsum_expr, x, initial_params['params']['kernel']) + 1.
     np.testing.assert_allclose(y, target, atol=1e-6)
 
-  @parameterized.parameters([((3,),), (3,)])
-  def test_conv(self, kernel_size):
+  def test_conv(self):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((1, 8, 3))
     conv_module = nn.Conv(
         features=4,
-        kernel_size=kernel_size,
+        kernel_size=(3,),
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -176,13 +175,12 @@ def test_conv(self, kernel_size):
     self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
     np.testing.assert_allclose(y, np.full((1, 6, 4), 10.))
 
-  @parameterized.parameters([((3,),), (3,)])
-  def test_single_input_conv(self, kernel_size):
+  def test_single_input_conv(self):
       rng = dict(params=random.PRNGKey(0))
       x = jnp.ones((8, 3))
       conv_module = nn.Conv(
           features=4,
-          kernel_size=kernel_size,
+          kernel_size=(3,),
           padding='VALID',
           kernel_init=initializers.ones,
           bias_init=initializers.ones,
@@ -191,13 +189,12 @@ def test_single_input_conv(self, kernel_size):
       self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
       np.testing.assert_allclose(y, np.full((6, 4), 10.))
 
-  @parameterized.parameters([((3,),), (3,)])
-  def test_group_conv(self, kernel_size):
+  def test_group_conv(self):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((1, 8, 4))
     conv_module = nn.Conv(
         features=4,
-        kernel_size=kernel_size,
+        kernel_size=(3,),
         feature_group_count=2,
         padding='VALID',
         kernel_init=initializers.ones,
@@ -207,13 +204,12 @@ def test_group_conv(self, kernel_size):
     self.assertEqual(initial_params['params']['kernel'].shape, (3, 2, 4))
     np.testing.assert_allclose(y, np.full((1, 6, 4), 7.))
 
-  @parameterized.parameters([((3,),), (3,)])
-  def test_conv_transpose(self, kernel_size):
+  def test_conv_transpose(self):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((1, 8, 3))
     conv_transpose_module = nn.ConvTranspose(
         features=4,
-        kernel_size=kernel_size,
+        kernel_size=(3,),
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -232,13 +228,12 @@ def test_conv_transpose(self, kernel_size):
                               [ 4.,  4.,  4.,  4.]]])
     np.testing.assert_allclose(y, correct_ans)
 
-  @parameterized.parameters([((3,),), (3,)])
-  def test_single_input_conv_transpose(self, kernel_size):
+  def test_single_input_conv_transpose(self):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((8, 3))
     conv_transpose_module = nn.ConvTranspose(
         features=4,
-        kernel_size=kernel_size,
+        kernel_size=(3,),
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -257,6 +252,12 @@ def test_single_input_conv_transpose(self, kernel_size):
                               [ 4.,  4.,  4.,  4.]])
     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))
+    with self.assertRaises(TypeError):
+      conv.init(random.PRNGKey(0), x)
+
   def test_embed(self):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.arange(4)[None]