4 ms·
I'm curious, what does paradigmatic JAX look like? Is there an equivalent of picoGPT [1] for JAX? [1] https://github.com/jaymody/picoGPT/blob/main/gpt2.py http
by Scene_Cast2 2y ago
I'm curious, what does paradigmatic JAX look like? Is there an equivalent of picoGPT [1] for JAX?
[1] https://github.com/jaymody/picoGPT/blob/main/gpt2.py https://github.com/jaymody/picoGPT/blob/main/gpt2.py
- jdeaton 2y agoyeah it looks exactly like that file but replace "import numpy as np" with "import jax.numpy as np" :)