Allow users to seed random sequence generations

[rlouf]

When integrating HF's GPT2 in #30 I used python.random.seed() and python.random.randint to generate a random seed value for jax.random.PRNGKey(). This method does not allow the user to seed the sequence generation, which is however necessary for reproducibilty.

We can either:

• Ask users to set the seed value globally;
• Pass PRNG keys explicitly in the code.