Move GeneratorType to types.py

Author: BirkhoffG

jax_dataloader/_modidx.py

@@ -118,6 +118,7 @@
                                                                                        'jax_dataloader/tests.py'),
                                       'jax_dataloader.tests.test_shuffle_reproducible': ( 'tests.html#test_shuffle_reproducible',
                                                                                           'jax_dataloader/tests.py')},
+            'jax_dataloader.types': {},
             'jax_dataloader.utils': { 'jax_dataloader.utils.Config': ('utils.html#config', 'jax_dataloader/utils.py'),
                                       'jax_dataloader.utils.Config.default': ('utils.html#config.default', 'jax_dataloader/utils.py'),
                                       'jax_dataloader.utils.Generator': ('utils.html#generator', 'jax_dataloader/utils.py'),

jax_dataloader/core.py

@@ -8,6 +8,7 @@
 from .utils import *
 from .datasets import *
 from .loaders import *
+from .types import *
 
 # %% auto 0
 __all__ = ['SUPPORTED_DATASETS', 'DataloaderBackends', 'get_backend_compatibilities', 'DataLoader']

jax_dataloader/loaders/base.py

@@ -3,7 +3,8 @@
 # %% ../../nbs/loader.base.ipynb 3
 from __future__ import print_function, division, annotations
 from ..imports import *
-from ..utils import Generator, GeneratorType
+from ..utils import Generator
+from ..types import GeneratorType
 
 # %% auto 0
 __all__ = ['BaseDataLoader']

jax_dataloader/loaders/jax.py

@@ -5,7 +5,8 @@
 from ..imports import *
 from ..datasets import ArrayDataset, JAXDataset
 from . import BaseDataLoader
-from ..utils import get_config, asnumpy, Generator, GeneratorType
+from ..utils import get_config, asnumpy, Generator
+from ..types import GeneratorType
 from ..tests import *
 import jax_dataloader as jdl
 from threading import Thread, Event

jax_dataloader/loaders/tensorflow.py

@@ -5,7 +5,8 @@
 from ..imports import *
 from . import BaseDataLoader
 from ..datasets import Dataset, ArrayDataset, JAXDataset
-from ..utils import check_tf_installed, get_config, Generator, GeneratorType
+from ..utils import check_tf_installed, get_config, Generator
+from ..types import GeneratorType
 from ..tests import *
 from jax.tree_util import tree_map
 import warnings

jax_dataloader/loaders/torch.py

@@ -5,7 +5,8 @@
 from ..imports import *
 from . import BaseDataLoader
 from ..datasets import Dataset, ArrayDataset, JAXDataset
-from ..utils import check_pytorch_installed, get_config, Generator, GeneratorType
+from ..utils import check_pytorch_installed, get_config, Generator
+from ..types import GeneratorType
 from ..tests import *
 from jax.tree_util import tree_map
 import warnings

jax_dataloader/types.py

@@ -0,0 +1,4 @@
+from .imports import *
+from .utils import Generator
+
+GeneratorType = Union[Generator, jax.Array, 'torch.Generator']

jax_dataloader/utils.py

@@ -7,8 +7,8 @@
 import collections
 
 # %% auto 0
-__all__ = ['GeneratorType', 'Config', 'get_config', 'manual_seed', 'check_pytorch_installed', 'has_pytorch_tensor',
-           'check_hf_installed', 'check_tf_installed', 'Generator', 'asnumpy']
+__all__ = ['Config', 'get_config', 'manual_seed', 'check_pytorch_installed', 'has_pytorch_tensor', 'check_hf_installed',
+           'check_tf_installed', 'Generator', 'asnumpy']
 
 # %% ../nbs/utils.ipynb 7
 @dataclass
@@ -67,10 +67,12 @@ def check_tf_installed():
 
 # %% ../nbs/utils.ipynb 21
 class Generator:
+    """A wrapper around JAX and PyTorch generators. This is used to generate random numbers in a reproducible way."""
+
     def __init__(
         self, 
         *, 
-        generator: jrand.Array | torch.Generator = None,
+        generator: jax.Array | torch.Generator = None, # Optional generator
     ):
         self._seed = None
         self._jax_generator = None
@@ -118,8 +120,6 @@ def torch_generator(self) -> torch.Generator:
             raise ValueError("Neither pytorch generator or seed is specified.")
         return self._torch_generator
 
-GeneratorType = Union[Generator, jax.Array, 'torch.Generator']
-
 # %% ../nbs/utils.ipynb 26
 def asnumpy(x) -> np.ndarray:
     if isinstance(x, np.ndarray):

nbs/core.ipynb

@@ -45,7 +45,8 @@
     "from jax_dataloader.imports import *\n",
     "from jax_dataloader.utils import *\n",
     "from jax_dataloader.datasets import *\n",
-    "from jax_dataloader.loaders import *"
+    "from jax_dataloader.loaders import *\n",
+    "from jax_dataloader.types import *"
    ]
   },
   {

nbs/loader.base.ipynb

@@ -39,7 +39,8 @@
     "#| export\n",
     "from __future__ import print_function, division, annotations\n",
     "from jax_dataloader.imports import *\n",
-    "from jax_dataloader.utils import Generator, GeneratorType"
+    "from jax_dataloader.utils import Generator\n",
+    "from jax_dataloader.types import GeneratorType"
    ]
   },
   {