Add docs to get_backend_compatibilities

BirkhoffG · Jan 30, 2024 · 7ea9ca5 · 7ea9ca5
1 parent 93b28ea
commit 7ea9ca5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/jax_dataloader/core.py b/jax_dataloader/core.py
@@ -62,7 +62,8 @@ def _check_backend_compatibility(ds, backend: str):
     return DataLoader(ds, backend=backend)
 
 # %% ../nbs/core.ipynb 8
-def get_backend_compatibilities() -> dict[str, list[type]]:
+def get_backend_compatibilities() -> dict[str, list[type]]: # { backend: [supported datasets] }
+    """Return list of supported dataloader backends for each dataset type"""
 
     ds = {
         JAXDataset: ArrayDataset(np.array([1,2,3])),

diff --git a/nbs/core.ipynb b/nbs/core.ipynb
@@ -145,7 +145,8 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "def get_backend_compatibilities() -> dict[str, list[type]]:\n",
+    "def get_backend_compatibilities() -> dict[str, list[type]]: # { backend: [supported datasets] }\n",
+    "    \"\"\"Return list of supported dataloader backends for each dataset type\"\"\"\n",
     "\n",
     "    ds = {\n",
     "        JAXDataset: ArrayDataset(np.array([1,2,3])),\n",