And, more importantly, things that cannot be expressed that way tend to not be a good fit for GPU computing anyway (independently of the language / framework you are using).
[0]: `array` is a shortcut here, JAX is not limited to operations on arrays.
I'll have to disagree with you a little bit here. SIMT model of GPUs are quiet a bit more expressive than the numpy's SIMD model. As an obvious example, you'll have to manually maintain a mask to implement if/else i.e. code path divergence in SIMD. GPUs automatically does this and many more to make your life easier. And frankly, I find it lot more easier to reason about what should happen to one data point than a bunch of them together.
An interesting article I read recently that has some relevance to this discussion. https://pharr.org/matt/blog/2018/04/18/ispc-origins
To be fair though, modern GPUs are pretty good at branching and latency hiding, while numpy-style code has poor data locality unless you have a magic compiler.
When you chain a sequence of vectorized operations on arrays, loop fusion would save you from allocating memory for each intermediate variable, and the round trip time of moving it from RAM to CPU multiple times. I don’t know how good JAX’s JITted loop fusion is on CPU, but I’ve been very very impressed by Julia.
Eg: I had some Numpy code that took hours (and needed terabyte RAM) that was very straightforward to code in Julia, and needed only a few GB to finish in a few seconds — on my laptop.
I want to be able to think in arrays, but to also not have to materialize the arrays as much as possible.
Having access to high performance explicit loops and ifs/masks allows one to focus on the hard parts of the algorithms, rather than on the purely incidental puzzle how to best avoid spending time in the Python runtime.