trying to get JAX and GPU acceleration working in docker (WSL2)

docker-compose build

Go away and make a cup of coffee...

docker-compose up

docker--compose down --remove-orphans

• jax-ml/jax#6340
• https://aesara.readthedocs.io/en/latest/tutorial/using_gpu.html?highlight=gpu#
• https://www.pymc-labs.io/blog-posts/the-quickest-migration-guide-ever-from-pymc3-to-pymc-v40/
• https://www.pymc-labs.io/blog-posts/pymc-stan-benchmark/