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