TODO Code structured to run without requiring JAX Code structured to run without Pytorch Code structured for numba JIT for parallel CPU