Investigate using JAX

[J08nY]

JAX while focused on machine learning is a versatile tool for computing on GPUs that has a numpy-compatible API. We could investigate using it for trace operations on GPUs.