add jax optional dependency

jjcmoon · jjcmoon · commit 79d4d07123af · 2025-09-29T11:11:11.000+02:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -22,6 +22,9 @@ classifiers = [
 torch = [
     "torch",
 ]
+jax = [
+    "jax",
+]
 
 
 [project.urls]
