chore: simplify imports

walln · Oct 2, 2024 · 24c4cb7 · 24c4cb7
1 parent 55ecb44
commit 24c4cb7
Show file tree

Hide file tree

Showing 10 changed files with 10 additions and 14 deletions.
diff --git a/src/loadax/__init__.py b/src/loadax/__init__.py
@@ -12,6 +12,7 @@
 from loadax.dataset.dataset import Dataset as Dataset
 from loadax.dataset.huggingface import HuggingFaceDataset as HuggingFaceDataset
 from loadax.dataset.partial_dataset import PartialDataset as PartialDataset
+from loadax.dataset.sampled_dataset import SampledDataset as SampledDataset
 from loadax.dataset.sharded_dataset import ShardedDataset as ShardedDataset
 from loadax.dataset.shuffled_dataset import Shuffleable as Shuffleable
 from loadax.dataset.simple import SimpleDataset as SimpleDataset
diff --git a/tests/dataset/test_batch_mapped.py b/tests/dataset/test_batch_mapped.py
@@ -1,8 +1,8 @@
 import jax.numpy as jnp
 import pytest
 
+from loadax import SimpleDataset
 from loadax.dataset.dataset import MappedBatchDataset
-from loadax.dataset.simple import SimpleDataset
 
 
 @pytest.fixture

diff --git a/tests/dataset/test_combined.py b/tests/dataset/test_combined.py
@@ -1,7 +1,6 @@
 import pytest
 
-from loadax.dataset.combined_dataset import CombinedDataset
-from loadax.dataset.simple import SimpleDataset
+from loadax import CombinedDataset, SimpleDataset
 
 
 @pytest.fixture

diff --git a/tests/dataset/test_huggingface_dataset.py b/tests/dataset/test_huggingface_dataset.py
@@ -1,7 +1,7 @@
 import pytest
 from datasets import Dataset as HFDataset
 
-from loadax.dataset.huggingface import HuggingFaceDataset
+from loadax import HuggingFaceDataset
 
 
 @pytest.fixture

diff --git a/tests/dataset/test_mapped.py b/tests/dataset/test_mapped.py
@@ -1,8 +1,8 @@
 import jax.numpy as jnp
 import pytest
 
+from loadax import SimpleDataset
 from loadax.dataset.dataset import MappedDataset
-from loadax.dataset.simple import SimpleDataset
 
 
 @pytest.fixture

diff --git a/tests/dataset/test_partial.py b/tests/dataset/test_partial.py
@@ -1,7 +1,6 @@
 import pytest
 
-from loadax.dataset.partial_dataset import PartialDataset
-from loadax.dataset.simple import SimpleDataset
+from loadax import PartialDataset, SimpleDataset
 
 
 @pytest.fixture

diff --git a/tests/dataset/test_sampled.py b/tests/dataset/test_sampled.py
@@ -1,8 +1,7 @@
 import jax
 import pytest
 
-from loadax.dataset.sampled_dataset import SampledDataset
-from loadax.dataset.simple import SimpleDataset
+from loadax import SampledDataset, SimpleDataset
 
 
 @pytest.fixture

diff --git a/tests/dataset/test_sharded_dataset.py b/tests/dataset/test_sharded_dataset.py
@@ -4,8 +4,7 @@
 
 import pytest
 
-from loadax.dataset.sharded_dataset import ShardedDataset
-from loadax.dataset.simple import SimpleDataset
+from loadax import ShardedDataset, SimpleDataset
 
 
 def compute_expected_boundaries(

diff --git a/tests/dataset/test_simple_dataset.py b/tests/dataset/test_simple_dataset.py
@@ -1,7 +1,7 @@
 import jax
 import pytest
 
-from loadax.dataset.simple import SimpleDataset
+from loadax import SimpleDataset
 
 
 @pytest.fixture

diff --git a/tests/test_dataloader.py b/tests/test_dataloader.py
@@ -1,7 +1,6 @@
 import pytest
 
-from loadax.dataloader.loader import Dataloader
-from loadax.dataset.simple import SimpleDataset
+from loadax import Dataloader, SimpleDataset
 
 
 @pytest.fixture