@parameterized.parameters([(((- 2), 3), (), 'bijk,jklm->bilm'), ((3, (- 2)), (), 'bijk,jklm->bilm'), (((- 2), 3), (0,), 'bijk,bjklm->bilm')])
def test_dense_general_vs_numpy(self, axis, batch_dims, einsum_expr):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((16, 8, 9, 10))
    dg_module = nn.DenseGeneral(features=(11, 12), axis=axis, batch_dims=batch_dims, bias_init=initializers.ones, kernel_init=initializers.normal())
    (y, initial_params) = dg_module.init_with_output(rng, x)
    target = (np.einsum(einsum_expr, x, initial_params['params']['kernel']) + 1.0)
    np.testing.assert_allclose(y, target, atol=1e-06)
@parameterized.parameters([(((- 2), 3), (), 'bijk,jklm->bilm'), ((3, (- 2)), (), 'bijk,jklm->bilm'), (((- 2), 3), (0,), 'bijk,bjklm->bilm')])
def test_dense_general_vs_numpy(self, axis, batch_dims, einsum_expr):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((16, 8, 9, 10))
    dg_module = nn.DenseGeneral(features=(11, 12), axis=axis, batch_dims=batch_dims, bias_init=initializers.ones, kernel_init=initializers.normal())
    (y, initial_params) = dg_module.init_with_output(rng, x)
    target = (np.einsum(einsum_expr, x, initial_params['params']['kernel']) + 1.0)
    np.testing.assert_allclose(y, target, atol=1e-06)