Possible typo in Conv module

[ozencgungor]

Hi there,

There seems to be a small typo in the docstring of the abstract _Conv module. It reads:

class _Conv(Module):
  """Convolution Module wrapping `lax.conv_general_dilated[_local]`.

  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: an integer or a sequence of `n` integers, representing the
      inter-window strides (default: 1).
....

But when I try to pass an integer to the kernel_size argument for a 1D Convolution acting on data of shape (N, M, F) where in channels last format as in the code example below:

###
        self.filter = nn.Conv(features=self.Fout, kernel_size=4**self.p, strides=4**self.p,
                              padding='VALID', use_bias=self.use_bias, 
                              kernel_init=self.kernel_init, bias_init=self.bias_init,
                              dtype=self.dtype, param_dtype=self.param_dtype)
    
    def __call__(self, inputs):
        x = jnp.take(inputs, self.masked_indices, axis=-2) #input format (N)MF
        x = self.filter(x)
        return x
###

I get the error The kernel size must be specified as a tuple/list of integers (eg.: [3, 3]).. Everything works fine if i set kernel_size=(4**self.p, ). I was wondering if this is a simple typo in the _Conv module docstring.

Thanks in advance.