def test_dense_extra_batch_dims(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((1, 2, 3))
    dense_module = nn.Dense(features=4, kernel_init=initializers.ones, bias_init=initializers.ones)
    (y, _) = dense_module.init_with_output(rng, x)
    np.testing.assert_allclose(y, np.full((1, 2, 4), 4.0))
def test_dense_extra_batch_dims(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((1, 2, 3))
    dense_module = nn.Dense(features=4, kernel_init=initializers.ones, bias_init=initializers.ones)
    (y, _) = dense_module.init_with_output(rng, x)
    np.testing.assert_allclose(y, np.full((1, 2, 4), 4.0))