Haiku & Jax weights initialisation

Question

In Pytorch the following code can be used to initialise a layer:

def init_layer(in_features, out_features):
 x = nn.Linear(in_features, out_features)
 limit = 1.0 / math.sqrt(in_features)
 x.weight = nn.Parameter(
    data=torch.distributions.uniform.Uniform(-limit, limit).sample(x.weight.shape), requires_grad=True
)
 return x

How to do the same thing using Jax & Haiku?

Thanks!