def test_dense_general_batch_dim_raises(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((1, 3, 2, 5))
    with self.assertRaises(ValueError):
        dg_module = nn.DenseGeneral(features=4, batch_dims=(0, 2), kernel_init=initializers.ones, bias_init=initializers.ones)
        dg_module.init_with_output(rng, x)
def test_dense_general_batch_dim_raises(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((1, 3, 2, 5))
    with self.assertRaises(ValueError):
        dg_module = nn.DenseGeneral(features=4, batch_dims=(0, 2), kernel_init=initializers.ones, bias_init=initializers.ones)
        dg_module.init_with_output(rng, x)