[Scene_Cast2]

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

[jdeaton]

yeah it looks exactly like that file but replace "import numpy as np" with "import jax.numpy as np" :)