If you have to express your computation using an "array programming DSL" with things like scan and gather anyways - why not opt to use torch/tensorflow/jax or anything else that targets MLIR? An example of writing a relu using an embedded array DSL is really not helping your case either - that's exactly the problem that these other solutions mentioned above are successfully solving for the past ~15y (starting with theano etc). Not sure what this brings to the table - doing that AoT instead of at runtime?
The goal of this work is to run existing unmodified CPU libraries (which may use core::simd) on the GPU. If you are manually writing ML-shaped workloads, it doesn't add any value over writing with tech like torch/tensorflow/jax which are custom built for those use-cases (except maybe familiarity if you are a CPU programmer).