From 8f45e16703e8bb283f486a144fddd48ff238e06b Mon Sep 17 00:00:00 2001
From: Tom Vercauteren <tom.vercauteren@kcl.ac.uk>
Date: Sun, 13 Nov 2022 17:58:20 +0000
Subject: [PATCH] list jax as optional dependency in setup.py

---
 setup.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/setup.py b/setup.py
index 0a039e8..21b0d8c 100644
--- a/setup.py
+++ b/setup.py
@@ -26,5 +26,8 @@ def readme():
     install_requires=[
         "torch",
     ],
+    extras_require={
+        "extras": ["jax"],
+    },
     zip_safe=False,
 )