GPU accelerated numerical ops on OCaml arrays via pyml + jax?

[mseri]

I think one alternative way can be to write the numerical code using Owl and then use owl-symbolic to export it to onnx and via pyml execute the onnx model on the appropriate hardware