Release v0.3.1

Version 0.3.0 release (#145)

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