def test_dense_general_batch_dim(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((2, 1, 3, 5))
    state = {'counter': 0.0}

    def _counter_init(rng, shape, dtype, state):
        del rng, dtype
        state['counter'] += 1.0
        return jnp.full(shape, state['counter'])
    counter_init = functools.partial(_counter_init, state=state)
    dg_module = nn.DenseGeneral(features=7, axis=(3, (- 2)), batch_dims=0, bias_init=initializers.ones, kernel_init=counter_init)
    (y, _) = dg_module.init_with_output(rng, x)
    target = np.concatenate([np.full((1, 1, 7), 16.0), np.full((1, 1, 7), 31.0)], axis=0)
    np.testing.assert_allclose(y, target)
def test_dense_general_batch_dim(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((2, 1, 3, 5))
    state = {'counter': 0.0}

    def _counter_init(rng, shape, dtype, state):
        del rng, dtype
        state['counter'] += 1.0
        return jnp.full(shape, state['counter'])
    counter_init = functools.partial(_counter_init, state=state)
    dg_module = nn.DenseGeneral(features=7, axis=(3, (- 2)), batch_dims=0, bias_init=initializers.ones, kernel_init=counter_init)
    (y, _) = dg_module.init_with_output(rng, x)
    target = np.concatenate([np.full((1, 1, 7), 16.0), np.full((1, 1, 7), 31.0)], axis=0)
    np.testing.assert_allclose(y, target)