One of the reasons I'm really excited about JAX is that I hope it will allow me to write fast Python code without worrying about these details.