def _counter_init(rng, shape, dtype, state):
    del rng, dtype
    state['counter'] += 1.0
    return jnp.full(shape, state['counter'])
def _counter_init(rng, shape, dtype, state):
    del rng, dtype
    state['counter'] += 1.0
    return jnp.full(shape, state['counter'])