list jax as optional dependency in setup.py

cai4cai · Nov 22, 2022 · 8f45e16 · 8f45e16
1 parent 1bc5a80
commit 8f45e16
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -26,5 +26,8 @@ def readme():
     install_requires=[
         "torch",
     ],
+    extras_require={
+        "extras": ["jax"],
+    },
     zip_safe=False,
 )