add jaxpm in setup

DifferentiableUniverseInitiative · Mar 5, 2024 · 08ebddd · 08ebddd
1 parent ee9d11a
commit 08ebddd
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -24,5 +24,6 @@
         "optax>=0.1.4",
         "scikit-learn>=1.2.0",
         "wheel",
+        "JaxPM @ git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git"
     ],
 )