Merge pull request #95 from mehdiataei/main

Get ready for 0.2.1 release. Fixed a bug running JAX on CPU.

Added setup for packaging