Repository: google/flax

Issue:
Helpful error when kernel_size is not an array in Conv

### Problem you have encountered:
If you try to init a Conv module by setting `kernel_size` to an int, you get an unhelpful error message.

### What you expected to happen:
Helpful error message that explains I should set `kernel_size` to an array.

### Logs, error messages, etc:
`TypeError: object of type 'int' has no len()`

### Steps to reproduce:

```python
from flax import nn
from jax import numpy as jnp, random

class CNN(nn.Module):

    def apply(self, x):
        x = nn.Conv(x, features=32, kernel_size=3)
        x = nn.relu(x)
        return x

cnn = CNN.init(random.PRNGKey(0), jnp.ones((1, 28, 28, 1)))
```




Discussion:
Good point! We should improve our error message here. (Or maybe we should allow simply passing in a single int? What do other frameworks do here?)
I think both [PyTorch](https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html) and [Tensorflow](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2D) allow passing a single int which means the same value will be used for all dimensions. 
Sure, then by all means let's also do that. We'll take a pull request for this (with tests) if anyone is interested.

Resolving patch:
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -195,7 +195,9 @@ class Conv(Module):
 
   Args:
     features: number of convolution filters.
-    kernel_size: shape of the convolutional kernel.
+    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.
     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
@@ -219,7 +221,7 @@ class Conv(Module):
     bias_init: initializer for the bias.
   """
   features: int
-  kernel_size: Sequence[int]
+  kernel_size: Union[int, Sequence[int]]
   strides: Optional[Sequence[int]] = None
   padding: Union[str, Sequence[Tuple[int, int]]] = 'SAME'
   input_dilation: Optional[Sequence[int]] = None
@@ -244,8 +246,13 @@ def __call__(self, inputs: Array) -> Array:
 
     inputs = jnp.asarray(inputs, self.dtype)
 
+    if isinstance(self.kernel_size, int):
+      kernel_size = (self.kernel_size,)
+    else:
+      kernel_size = self.kernel_size
+
     is_single_input = False
-    if inputs.ndim == len(self.kernel_size) + 1:
+    if inputs.ndim == len(kernel_size) + 1:
       is_single_input = True
       inputs = jnp.expand_dims(inputs, axis=0)
 
@@ -254,7 +261,7 @@ def __call__(self, inputs: Array) -> Array:
 
     in_features = inputs.shape[-1]
     assert in_features % self.feature_group_count == 0
-    kernel_shape = self.kernel_size + (
+    kernel_shape = kernel_size + (
         in_features // self.feature_group_count, self.features)
     kernel = self.param('kernel', self.kernel_init, kernel_shape)
     kernel = jnp.asarray(kernel, self.dtype)
@@ -285,7 +292,9 @@ class ConvTranspose(Module):
 
   Args:
     features: number of convolution filters.
-    kernel_size: shape of the convolutional kernel.
+    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.
     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
@@ -303,7 +312,7 @@ class ConvTranspose(Module):
     bias_init: initializer for the bias.
   """
   features: int
-  kernel_size: Sequence[int]
+  kernel_size: Union[int, Sequence[int]]
   strides: Optional[Sequence[int]] = None
   padding: Union[str, Sequence[Tuple[int, int]]] = 'SAME'
   kernel_dilation: Optional[Sequence[int]] = None
@@ -325,15 +334,21 @@ def __call__(self, inputs: Array) -> Array:
       The convolved data.
     """
     inputs = jnp.asarray(inputs, self.dtype)
+
+    if isinstance(self.kernel_size, int):
+      kernel_size = (self.kernel_size,)
+    else:
+      kernel_size = self.kernel_size
+
     is_single_input = False
-    if inputs.ndim == len(self.kernel_size) + 1:
+    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)
 
     in_features = inputs.shape[-1]
-    kernel_shape = self.kernel_size + (in_features, self.features)
+    kernel_shape = kernel_size + (in_features, self.features)
     kernel = self.param('kernel', self.kernel_init, kernel_shape)
     kernel = jnp.asarray(kernel, self.dtype)
 
diff --git a/flax/nn/linear.py b/flax/nn/linear.py
--- a/flax/nn/linear.py
+++ b/flax/nn/linear.py
@@ -192,7 +192,9 @@ def apply(self,
     Args:
       inputs: input data with dimensions (batch, spatial_dims..., features).
       features: number of convolution filters.
-      kernel_size: shape of the convolutional kernel.
+      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.
       padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
@@ -219,12 +221,14 @@ def apply(self,
     """
 
     inputs = jnp.asarray(inputs, dtype)
+    if isinstance(kernel_size, int):
+      kernel_size = (kernel_size,)
 
     is_single_input = False
     if inputs.ndim == len(kernel_size) + 1:
       is_single_input = True
       inputs = jnp.expand_dims(inputs, axis=0)
-      
+
     if strides is None:
       strides = (1,) * (inputs.ndim - 2)
 
@@ -276,7 +280,9 @@ def apply(self,
     Args:
       inputs: input data with dimensions (batch, spatial_dims..., features).
       features: number of convolution filters.
-      kernel_size: shape of the convolutional kernel.
+      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.
       padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
@@ -296,12 +302,14 @@ def apply(self,
       The convolved data.
     """
     inputs = jnp.asarray(inputs, dtype)
+    if isinstance(kernel_size, int):
+      kernel_size = (kernel_size,)
     
     is_single_input = False
     if inputs.ndim == len(kernel_size) + 1:
       is_single_input = True
       inputs = jnp.expand_dims(inputs, axis=0)
-    
+
     strides = strides or (1,) * (inputs.ndim - 2)
 
     in_features = inputs.shape[-1]


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
@@ -163,12 +163,13 @@ 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)
 
-  def test_conv(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_conv(self, kernel_size):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((1, 8, 3))
     conv_module = nn.Conv(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -177,12 +178,13 @@ def test_conv(self):
     self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
     np.testing.assert_allclose(y, np.full((1, 6, 4), 10.))
 
-  def test_single_input_conv(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_single_input_conv(self, kernel_size):
       rng = dict(params=random.PRNGKey(0))
       x = jnp.ones((8, 3))
       conv_module = nn.Conv(
           features=4,
-          kernel_size=(3,),
+          kernel_size=kernel_size,
           padding='VALID',
           kernel_init=initializers.ones,
           bias_init=initializers.ones,
@@ -191,12 +193,13 @@ def test_single_input_conv(self):
       self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
       np.testing.assert_allclose(y, np.full((6, 4), 10.))
 
-  def test_group_conv(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_group_conv(self, kernel_size):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((1, 8, 4))
     conv_module = nn.Conv(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         feature_group_count=2,
         padding='VALID',
         kernel_init=initializers.ones,
@@ -206,12 +209,13 @@ 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.))
 
-  def test_conv_transpose(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_conv_transpose(self, kernel_size):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((1, 8, 3))
     conv_transpose_module = nn.ConvTranspose(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -230,12 +234,13 @@ def test_conv_transpose(self):
                               [ 4.,  4.,  4.,  4.]]])
     np.testing.assert_allclose(y, correct_ans)
 
-  def test_single_input_conv_transpose(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_single_input_conv_transpose(self, kernel_size):
     rng = dict(params=random.PRNGKey(0))
     x = jnp.ones((8, 3))
     conv_transpose_module = nn.ConvTranspose(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
diff --git a/tests/nn_linear_test.py b/tests/nn_linear_test.py
--- a/tests/nn_linear_test.py
+++ b/tests/nn_linear_test.py
@@ -162,12 +162,13 @@ def test_dense_general_vs_numpy(self, axis, batch_dims, einsum_expr):
     target = onp.einsum(einsum_expr, x, dg_module.params['kernel']) + 1.
     onp.testing.assert_allclose(y, target, atol=1e-6)
 
-  def test_conv(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_conv(self, kernel_size):
     rng = random.PRNGKey(0)
     x = jnp.ones((1, 8, 3))
     conv_module = nn.Conv.partial(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -176,13 +177,14 @@ def test_conv(self):
     model = nn.Model(conv_module, initial_params)
     self.assertEqual(model.params['kernel'].shape, (3, 3, 4))
     onp.testing.assert_allclose(y, onp.full((1, 6, 4), 10.))
-    
-  def test_single_input_conv(self):
+
+  @parameterized.parameters([((3,),), (3,)])
+  def test_single_input_conv(self, kernel_size):
     rng = random.PRNGKey(0)
     x = jnp.ones((8, 3))
     conv_module = nn.Conv.partial(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -192,12 +194,13 @@ def test_single_input_conv(self):
     self.assertEqual(model.params['kernel'].shape, (3, 3, 4))
     onp.testing.assert_allclose(y, onp.full((6, 4), 10.))
 
-  def test_group_conv(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_group_conv(self, kernel_size):
     rng = random.PRNGKey(0)
     x = jnp.ones((1, 8, 4))
     conv_module = nn.Conv.partial(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         feature_group_count=2,
         padding='VALID',
         kernel_init=initializers.ones,
@@ -208,12 +211,13 @@ def test_group_conv(self):
     self.assertEqual(model.params['kernel'].shape, (3, 2, 4))
     onp.testing.assert_allclose(y, onp.full((1, 6, 4), 7.))
 
-  def test_conv_transpose(self):
+  @parameterized.parameters([((3,),), (3,)])
+  def test_conv_transpose(self, kernel_size):
     rng = random.PRNGKey(0)
     x = jnp.ones((1, 8, 3))
     conv_transpose_module = nn.ConvTranspose.partial(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,
@@ -232,13 +236,14 @@ def test_conv_transpose(self):
                               [ 7.,  7.,  7.,  7.],
                               [ 4.,  4.,  4.,  4.]]])
     onp.testing.assert_allclose(y, correct_ans)
-    
-  def test_single_input_conv_transpose(self):
+
+  @parameterized.parameters([((3,),), (3,)])
+  def test_single_input_conv_transpose(self, kernel_size):
     rng = random.PRNGKey(0)
     x = jnp.ones((8, 3))
     conv_transpose_module = nn.ConvTranspose.partial(
         features=4,
-        kernel_size=(3,),
+        kernel_size=kernel_size,
         padding='VALID',
         kernel_init=initializers.ones,
         bias_init=initializers.ones,