search

jaymody / picoGPT

An unnecessarily tiny implementation of GPT-2 in NumPy.
MIT License
3.25k stars 417 forks source link

add full support for JAX in utils.py #17

Open aalmarhabi opened 1 year ago

aalmarhabi commented 1 year ago

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 commented 1 year ago

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