Comment by nivter

If it's about teaching and not about efficiency, why not just use plain Python? One could argue it is actually better since students don't have to worry about typing and syntax, and it allows a gentler introduction to commonly used tools like jax and numpy while getting comfortable with the language.