def test_dense_is_dense_general(self):
    x = jax.random.normal(random.PRNGKey(0), (5, 3))
    dense_module = nn.Dense(features=4, use_bias=True, bias_init=initializers.normal())
    (y1, _) = dense_module.init_with_output(dict(params=random.PRNGKey(1)), x)
    dg_module = nn.DenseGeneral(features=4, use_bias=True, bias_init=initializers.normal())
    (y2, _) = dg_module.init_with_output(dict(params=random.PRNGKey(1)), x)
    np.testing.assert_allclose(y1, y2)
def test_dense_is_dense_general(self):
    x = jax.random.normal(random.PRNGKey(0), (5, 3))
    dense_module = nn.Dense(features=4, use_bias=True, bias_init=initializers.normal())
    (y1, _) = dense_module.init_with_output(dict(params=random.PRNGKey(1)), x)
    dg_module = nn.DenseGeneral(features=4, use_bias=True, bias_init=initializers.normal())
    (y2, _) = dg_module.init_with_output(dict(params=random.PRNGKey(1)), x)
    np.testing.assert_allclose(y1, y2)