Allow for cuda support in pyproject

Author: plainerman

README.md

@@ -1,7 +1,7 @@
 <h1 align="center">Doob’s Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling</h1>
 <p align="center">
 <a href="https://github.com/plainerman/variational-doob"><img src="https://img.shields.io/badge/python-3670A0?style=for-the-badge&logo=python&logoColor=ffdd54" alt="Built with Python"/></a>
-<a href="https://github.com/plainerman/Variational-Doob/blob/main/notebooks/tps_gaussian.ipynb"><img src="https://img.shields.io/badge/jupyter-%23FA0F00.svg?style=for-the-badge&logo=jupyter&logoColor=white" alt="Jupyter"/></a>
+<a href="https://github.com/plainerman/variational-doob/blob/main/notebooks/tps_gaussian.ipynb"><img src="https://img.shields.io/badge/jupyter-%23FA0F00.svg?style=for-the-badge&logo=jupyter&logoColor=white" alt="Jupyter"/></a>
 <a href="https://github.com/jax-ml/jax"><img src="https://img.shields.io/badge/library-JAX-5f0964?style=for-the-badge" alt="Jax"/></a>
 </p>
 <p align="center">

pyproject.toml

@@ -12,7 +12,8 @@ dependencies = [
     "scipy==1.12.0",
     "scikit-image>=0.24.0",
     "openpathsampling>=1.6.1",
-    "jax==0.4.23",
+    "jax==0.4.23; platform_system == 'Darwin'",
+    "jax[cuda12]==0.4.23; platform_system != 'Darwin'",
     "jaxlib==0.4.23",
     "flax==0.8.3",
     "notebook>=7.2.2",