Repository: google/flax

Issue:
NNX Conv documentation states that kernel size can be an integer for 1D convs.
https://github.com/google/flax/blob/829e004c6d5646f24bd048023a4407e3b4d5aa98/flax/experimental/nnx/nnx/nn/linear.py#L368

The NNX Conv documentation says that you can pass kernel size as an integer for 1D convolutions, however, the very first thing the constructor does is throw a TypeError if the kernel size is an integer.


Discussion:

Resolving patch:
diff --git a/flax/experimental/nnx/nnx/nn/linear.py b/flax/experimental/nnx/nnx/nn/linear.py
--- a/flax/experimental/nnx/nnx/nn/linear.py
+++ b/flax/experimental/nnx/nnx/nn/linear.py
@@ -371,8 +371,9 @@ class Conv(Module):
   Attributes:
     features: number of convolution filters.
     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.
+      the kernel size can be passed as an integer, which will be interpreted
+      as a tuple of the single integer. For all other cases, it must 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'`, the string
@@ -407,7 +408,7 @@ def __init__(
     self,
     in_features: int,
     out_features: int,
-    kernel_size: tp.Sequence[int],
+    kernel_size: int | tp.Sequence[int],
     strides: tp.Union[None, int, tp.Sequence[int]] = 1,
     *,
     padding: PaddingLike = 'SAME',
@@ -429,11 +430,7 @@ def __init__(
     rngs: rnglib.Rngs,
   ):
     if isinstance(kernel_size, int):
-      raise TypeError(
-        'Expected Conv kernel_size to be a'
-        ' tuple/list of integers (eg.: [3, 3]) but got'
-        f' {kernel_size}.'
-      )
+      kernel_size = (kernel_size,)
     else:
       kernel_size = tuple(kernel_size)
 
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -330,8 +330,9 @@ class _Conv(Module):
 
   Attributes:
     features: number of convolution filters.
-    kernel_size: shape of the convolutional kernel.
-    strides: an integer or a sequence of ``n`` integers, representing the
+    kernel_size: shape of the convolutional kernel. An integer will be
+      interpreted as a tuple of the single integer.
+    strides: an integer or a sequence of `n` integers, representing the
       inter-window strides (default: 1).
     padding: either the string ``'SAME'``, the string ``'VALID'``, the string
       ``'CIRCULAR'`` (periodic boundary conditions), or a sequence of ``n`` ``(low,
@@ -362,7 +363,7 @@ class _Conv(Module):
   """
 
   features: int
-  kernel_size: Sequence[int]
+  kernel_size: Union[int, Sequence[int]]
   strides: Union[None, int, Sequence[int]] = 1
   padding: PaddingLike = 'SAME'
   input_dilation: Union[None, int, Sequence[int]] = 1
@@ -414,12 +415,9 @@ def __call__(self, inputs: Array) -> Array:
       The convolved data.
     """
 
+    kernel_size: Sequence[int]
     if isinstance(self.kernel_size, int):
-      raise TypeError(
-        'Expected Conv kernel_size to be a'
-        ' tuple/list of integers (eg.: [3, 3]) but got'
-        f' {self.kernel_size}.'
-      )
+      kernel_size = (self.kernel_size,)
     else:
       kernel_size = tuple(self.kernel_size)
 
@@ -614,8 +612,9 @@ class Conv(_Conv):
 
   Attributes:
     features: number of convolution filters.
-    kernel_size: shape of the convolutional kernel.
-    strides: an integer or a sequence of ``n`` integers, representing the
+    kernel_size: shape of the convolutional kernel. An integer will be
+      interpreted as a tuple of the single integer.
+    strides: an integer or a sequence of `n` integers, representing the
       inter-window strides (default: 1).
     padding: either the string ``'SAME'``, the string ``'VALID'``, the string
       ``'CIRCULAR'`` (periodic boundary conditions), or a sequence of ``n`` ``(low,
@@ -679,8 +678,9 @@ class ConvLocal(_Conv):
 
   Attributes:
     features: number of convolution filters.
-    kernel_size: shape of the convolutional kernel.
-    strides: an integer or a sequence of ``n`` integers, representing the
+    kernel_size: shape of the convolutional kernel. An integer will be
+      interpreted as a tuple of the single integer.
+    strides: an integer or a sequence of `n` integers, representing the
       inter-window strides (default: 1).
     padding: either the string ``'SAME'``, the string ``'VALID'``, the string
       ``'CIRCULAR'`` (periodic boundary conditions), or a sequence of ``n`` ``(low,
@@ -745,12 +745,13 @@ class ConvTranspose(Module):
   Attributes:
     features: number of convolution filters.
     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'``, 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
+      the kernel size can be passed as an integer, which will be interpreted
+      as a tuple of the single 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'`, 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. A single int is interpreted as applying the same padding
       in all dims and assign a single int in a sequence causes the same padding
       to be used on both sides.


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
@@ -989,12 +989,12 @@ def test_transpose_kernel_conv_transpose(self, use_bias):
     self.assertEqual(initial_params['params']['kernel'].shape, (6, 6, 4, 3))
     self.assertEqual(y.shape, (1, 30, 30, 4))
 
-  @parameterized.product(module=(nn.Conv, nn.ConvLocal))
-  def test_int_kernel_size(self, module):
-    conv = module(features=4, kernel_size=3)
+  @parameterized.product(module=(nn.Conv, nn.ConvLocal, nn.ConvTranspose))
+  def test_int_kernel_equality(self, module):
+    conv_int = module(features=4, kernel_size=3)
+    conv_seq = module(features=4, kernel_size=(3,))
     x = jnp.ones((8, 3))
-    with self.assertRaises(TypeError):
-      conv.init(random.key(0), x)
+    self.assertTrue(jax.tree_util.tree_all(jax.tree_map(lambda x, y: (x==y).all(), conv_int.init(random.key(0), x), conv_seq.init(random.key(0), x))))
 
   def test_embed(self):
     rng = dict(params=random.key(0))