feat: add a to argument

.github/workflows/test-gpu.yaml

@@ -41,7 +41,7 @@ jobs:
     runs-on: "cirun-aws-gpu--${{ github.run_id }}"
     strategy:
       matrix:
-        extras: ["torch,cupy-cuda12", "torch", "cupy-cuda12"]
+        extras: ["torch,cupy-cuda12", "torch", "cupy-cuda12", "jax-cuda12"]
     # Setting a timeout of 30 minutes, as the AWS costs money
     # At time of writing, a typical run takes about 5 minutes
     timeout-minutes: 30

README.md

@@ -131,7 +131,7 @@ with ad.settings.override(remove_unused_categories=False):
         batch_size=4096,
         chunk_size=32,
         preload_nchunks=256,
-        to_torch=True
+        to="torch"
     )
     # `use_collection` automatically uses the on-disk `X` and full `obs` in the `Loader`
     # but the `load_adata` arg can override this behavior

docs/notebooks/example.ipynb

@@ -261,7 +261,7 @@
         "    preload_nchunks=32,  # Number of chunks to preload + shuffle - default settings should work well\n",
         "    # If True, preloaded chunks are moved to GPU memory via `cupy`, which can put more pressure on GPU memory but will accelerate loading ~20%\n",
         "    preload_to_gpu=False,\n",
-        "    to_torch=True,\n",
+        "    to=\"torch\",\n",
         ")\n",
         "\n",
         "# Add in the shuffled data that should be used for training.\n",

pyproject.toml

@@ -62,6 +62,13 @@ optional-dependencies.doc = [
   "sphinxcontrib-bibtex>=1",
   "sphinxext-opengraph",
 ]
+optional-dependencies.jax = [ "jax" ]
+optional-dependencies.jax-cuda12 = [
+  "jax[cuda12]",
+]
+optional-dependencies.jax-cuda13 = [
+  "jax[cuda13]",
+]
 optional-dependencies.test = [
   "annbatch[zarrs]",
   "coverage",
@@ -88,7 +95,7 @@ envs.docs.scripts.clean = "git clean -fdX -- {args:docs}"
 envs.hatch-test.features = [ "test" ]
 envs.hatch-test.python = "3.14"
 envs.hatch-test.matrix = [
-  { deps = [ "min-low", "pre", "torch", "min-high" ] },
+  { deps = [ "min-low", "pre", "torch", "min-high", "jax" ] },
 ]
 # If the matrix variable `deps` is set to "pre",
 # set the environment variable `UV_PRERELEASE` to "allow".
@@ -106,6 +113,7 @@ envs.hatch-test.overrides.matrix.deps.python = [
 ]
 envs.hatch-test.overrides.matrix.deps.features = [
   { if = [ "torch" ], value = "torch" },
+  { if = [ "jax" ], value = "jax" },
 ]
 
 [tool.ruff]

src/annbatch/compat.py

@@ -19,3 +19,10 @@
     from cupyx.scipy.sparse import csr_matrix as CupyCSRMatrix
 else:
     CupyCSRMatrix = type("csr_matrix", (), {"__module__": "cupyx.scipy.sparse"})
+
+if TYPE_CHECKING or find_spec("jax"):
+    from jax import Array as JaxArray
+    from jax.experimental.sparse import CSR as JAXCsrMatrix
+else:
+    JAXCsrMatrix = type("CSR", (), {"__module__": "jax.experimental.sparse"})
+    JaxArray = type("Array", (), {"__module__": "jax"})