Enable GPU acceleration

[NoRaincheck]

Candidate approaches:

• Use autograd successor: https://github.com/google/jax
• Use tensorflow

We might just wait for jax to be pip installable and support windows first; or have a wrapper that imports jax where available otherwise defaults to autograd