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:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/google/flax/pull/3537?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`50cd169`)](https://app.codecov.io/gh/google/flax/commit/50cd169debcc02c10926cac1f2264acfc38ebf20?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 56.03% compared to head [(`cd5a21f`)](https://app.codecov.io/gh/google/flax/pull/3537?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 56.03%.

| [Files](https://app.codecov.io/gh/google/flax/pull/3537?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) | Patch % | Lines |
|---|---|---|
| [flax/experimental/nnx/nnx/nn/linear.py](https://app.codecov.io/gh/google/flax/pull/3537?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google#diff-ZmxheC9leHBlcmltZW50YWwvbm54L25ueC9ubi9saW5lYXIucHk=) | 0.00% | [1 Missing :warning: ](https://app.codecov.io/gh/google/flax/pull/3537?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) |

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #3537   +/-   ##
=======================================
  Coverage   56.03%   56.03%           
=======================================
  Files         100      100           
  Lines       11867    11867           
=======================================
  Hits         6650     6650           
  Misses       5217     5217           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/google/flax/pull/3537?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).


Patch:
diff --git a/flax/experimental/nnx/nnx/nn/linear.py b/flax/experimental/nnx/nnx/nn/linear.py
index 6308f579..1c0c0afc 100644
--- 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 @@ class Conv(Module):
     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 @@ class Conv(Module):
     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
index b60fb2af..577b1d31 100644
--- 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 @@ class _Conv(Module):
       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
index 1399f56a..e46f85fc 100644
--- a/tests/linen/linen_linear_test.py
+++ b/tests/linen/linen_linear_test.py
@@ -989,12 +989,12 @@ class LinearTest(parameterized.TestCase):
     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))