workaround jaxlib 0.3 removal from pypi

graphcore-research · Nov 10, 2023 · 3d3e26e · 3d3e26e
1 parent d751303
commit 3d3e26e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/requirements_cpu.txt b/requirements_cpu.txt
@@ -6,4 +6,4 @@
 #    requirements_test.txt for test-only dependencies    
 
 jax==0.3.16
-jaxlib==0.3.15
+jaxlib @ https://storage.googleapis.com/jax-releases/nocuda/jaxlib-0.3.15-cp38-none-manylinux2014_x86_64.whl
diff --git a/setup.py b/setup.py
@@ -41,8 +41,8 @@ def remove_comments(line: str):
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python :: 3",
     ],
-    install_requires=install_requires + cpu_requires,
-    extras_require={"ipu": ipu_requires, "test": test_requires},
+    install_requires=install_requires,
+    extras_require={"cpu": cpu_requires, "ipu": ipu_requires, "test": test_requires},
     python_requires=">=3.8",
     packages=["pyscf_ipu"],
     entry_points={"console_scripts": ["nanoDFT=pyscf_ipu.nanoDFT:main"]},