export alias

jaxonloader/__init__.py

@@ -1,5 +1,8 @@
 from jaxonloader._datasets import *  # noqa
 from jaxonloader.dataloader import JaxonDataLoader, make  # noqa
 import equinox as eqx
+from jaxtyping import Array
+from collections.abc import Callable
 
 Index = eqx.nn.State
+JITJaxonDataLoader = Callable[[eqx.nn.State], tuple[Array, eqx.nn.State, bool]]

jaxonloader/dataloader.py

@@ -1,10 +1,9 @@
-from collections.abc import Callable
-
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, PRNGKeyArray
 
+from jaxonloader import JITJaxonDataLoader
 from jaxonloader.dataset import JaxonDataset
 
 
@@ -78,10 +77,7 @@ def make(
     drop_last: bool = False,
     key: PRNGKeyArray | None = None,
     jit: bool = True,
-) -> (
-    tuple[Callable[[eqx.nn.State], tuple[Array, eqx.nn.State, bool]], eqx.nn.State]
-    | tuple[JaxonDataLoader, eqx.nn.State]
-):
+) -> tuple[JITJaxonDataLoader, eqx.nn.State] | tuple[JaxonDataLoader, eqx.nn.State]:
     dataloader, index = eqx.nn.make_with_state(JaxonDataLoader)(
         dataset, batch_size, shuffle, drop_last, key=key
     )

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxonloader"
-version = "0.2.1"
+version = "0.2.2"
 description = "A dataloader, but for JAX"
 readme = "README.md"
 requires-python ="~=3.10"