Is NVIDIA's JIT-based approach here similar JAX's, except targeting CUDA directly rather than XLA? Would like to know how these different JIT compilers relate to one another.