Adam and He initialisation

Implement the most used optimiser in deep learning and the standard ReLU initialisation.

• he_init(n_in, n_out, rng): return weights of shape (n_in, n_out) drawn from a normal distribution with mean 0 and standard deviation $\sqrt{2/n_{\text{in}}}$, using rng.normal.
• adam_step(param, grad, m, v, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8): one Adam update as in the lesson, for step number t (starting at 1). Return the new (param, m, v). Don't modify the inputs in place.