Jax basically gives you the array language without leaving Python, and the XLA backend means you're not hand-tuning C for the GPU path. The numpy interface limi

Jax basically gives you the array language without leaving Python, and the XLA backend means you're not hand-tuning C for the GPU path. The numpy interface limitation is real though and once you need something that doesn't map cleanly to vectorized ops, you're either fighting the abstraction or dropping down anyway. The daxpby example is a good one. Every time BLAS adds another special-case routine it's basically admitting the interface wasn't general enough. At some point you're just writing C with extra steps.