JAX can’t do what Numba can do for example. I just want one way to write simple math-y code like you normally would and automagically convert to run on one of the above approaches.

That’s what compilers and high level languages are supposed to be for!