add full support for JAX in utils.py

[aalmarhabi]

Hi Jay, thank you for making the PicoGPT repo. It is very useful, as well as the blog I was following it to understand the GPT model more.

The current version of the repo does support the JAX numpy API. However, for full support, you may consider these changes. I found out that if I also change utils.py numpy to jax.numpy an issue arises (Issue 4564).

So if you modified both gpt2.py and utlis.py to use jax.numpy on the new changes, it should work fine. Also for the current version work fine.

in gpt2 function token + positional embeddings

    x = wte[np.array(inputs)] + wpe[np.array(range(len(inputs)))]  # [n_seq] -> [n_seq, n_embd]

[aalmarhabi]

Also, I was testing the code on Python 3.10.9 version, I just forgot to mention that.