@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=kernel_size, padding='VALID', kernel_init=initializers.ones, bias_init=initializers.ones)
    (y, initial_params) = conv_module.init_with_output(rng, x)
    self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
    np.testing.assert_allclose(y, np.full((6, 4), 10.0))
@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=kernel_size, padding='VALID', kernel_init=initializers.ones, bias_init=initializers.ones)
    (y, initial_params) = conv_module.init_with_output(rng, x)
    self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
    np.testing.assert_allclose(y, np.full((6, 4), 10.0))