In addition to what @fehiepsi suggested can you check the following thread and run the example there to see if that’s work on your machine. The versions of numpyro and dill packages are also provided at the end of the thread.
numpyro
dill