Backport PR #356: ➕ dep-add(extras): add extra dependencies (#357)

Co-authored-by: Nathaniel Starkman <[email protected]>

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ default_stages: [pre-commit, pre-push]
 
 repos:
   - repo: https://github.com/commitizen-tools/commitizen
-    rev: v4.0.0
+    rev: v4.1.0
     hooks:
       - id: commitizen
         additional_dependencies: [cz-conventional-gitmoji]
@@ -50,7 +50,7 @@ repos:
       - id: check-readthedocs
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.8.1"
+    rev: "v0.8.6"
     hooks:
       # Run the linter
       - id: ruff
@@ -67,14 +67,14 @@ repos:
         additional_dependencies: [black==23.*]
 
   - repo: https://github.com/rbubley/mirrors-prettier
-    rev: "v3.4.1"
+    rev: "v3.4.2"
     hooks:
       - id: prettier
         types_or: [yaml, markdown, html, css, scss, javascript, json]
         args: [--prose-wrap=always]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v1.13.0"
+    rev: "v1.14.1"
     hooks:
       - id: mypy
         files: src
@@ -86,11 +86,11 @@ repos:
     hooks:
       - id: codespell
 
-  - repo: https://github.com/abravalheri/validate-pyproject
-    rev: v0.23
-    hooks:
-      - id: validate-pyproject
-        additional_dependencies: ["validate-pyproject-schema-store[all]"]
+  # - repo: https://github.com/abravalheri/validate-pyproject
+  #   rev: v0.23
+  #   hooks:
+  #     - id: validate-pyproject
+  #       additional_dependencies: ["validate-pyproject-schema-store[all]"]
 
   - repo: local
     hooks:

pyproject.toml

@@ -43,6 +43,14 @@
   # NOTE: zeroth is only required for matplotlib
 
   [project.optional-dependencies]
+    # jax extras
+    cpu             = ["jax[cpu]"]
+    cuda            = ["jax[cuda]"]
+    cuda12          = ["jax[cuda12]"]
+    cuda12_local    = ["jax[cuda12_local]"]
+    rocm            = ["jax[rocm]"]
+    k8s             = ["jax[k8s]"]
+    # extra extras
     all             = ["unxt[backend-astropy,interop-gala,interop-mpl]"]
     backend-astropy = ["astropy>=6.0"]
     interop-gala    = ["gala>=1.8"]
@@ -261,6 +269,15 @@ UNXT_ENABLE_RUNTIME_TYPECHECKING = "beartype.beartype"
 
 
 [tool.uv]
+# Require that `unxt[<jax-extras>]` requirements are resolved in different forks
+# so that they cannot conflict with one another.
+conflicts = [
+    [
+        { extra = "cpu" },
+        { extra = "cuda" }, { extra = "cuda12" }, { extra = "cuda12_local" },
+        { extra = "rocm" }, { extra = "k8s" }
+    ]
+]
 constraint-dependencies = [
   "jax<0.4.36", # TODO: remove when https://github.com/patrick-kidger/quax/pull/37 is merged
 ]