@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=kernel_size, feature_group_count=2, 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, 2, 4))
    np.testing.assert_allclose(y, np.full((1, 6, 4), 7.0))
@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=kernel_size, feature_group_count=2, 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, 2, 4))
    np.testing.assert_allclose(y, np.full((1, 6, 4), 7.0))