An alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can

An alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can interoperate with the larger ecosystem, while still being fast to run.