shard_map and random generators

[PhilipVinc]

If I do

In [22]: mesh =jax.make_mesh((4,), ('S',), axis_types=(AxisType.Explicit,))
In [23]: jax.sharding.set_mesh(mesh)

In [26]: @jax.shard_map(out_specs=jax.P('S'))
    ...: def r(k):
    ...:     print(k)
    ...:     return jax.random.normal(k, (1,))

In [27]: r(k)
Out[27]: Array([-1.18428442, -1.18428442, -1.18428442, -1.18428442], dtype=float64)

so I wonder what's the correct way to ensure that this output matches

In [28]: jax.random.normal(k, (4,), out_sharding=jax.P('S'))
Out[28]: Array([-1.18428442, -0.11617041,  0.17269028,  0.95730718], dtype=float64)

should I generate a vector of keys before calling shard map? or should I use a jax.random.fold_in inside the shard map context?
what is done by th standard rngs in jax?