Update test cases for testing num of batches

jax_dataloader/tests.py

@@ -18,6 +18,7 @@ def get_batch(batch):
 # %% ../nbs/tests.ipynb 4
 def test_no_shuffle(cls, ds, batch_size: int, feats, labels):
     dl = cls(ds, batch_size=batch_size, shuffle=False)
+    assert len(dl) == len(feats) // batch_size + 1
     for _ in range(2):
         X_list, Y_list = [], []
         for batch in dl:
@@ -31,6 +32,7 @@ def test_no_shuffle(cls, ds, batch_size: int, feats, labels):
 # %% ../nbs/tests.ipynb 5
 def test_no_shuffle_drop_last(cls, ds, batch_size: int, feats, labels):
     dl = cls(ds, batch_size=batch_size, shuffle=False, drop_last=True)
+    assert len(dl) == len(feats) // batch_size
     for _ in range(2):
         X_list, Y_list = [], []
         for batch in dl:
@@ -46,6 +48,7 @@ def test_no_shuffle_drop_last(cls, ds, batch_size: int, feats, labels):
 def test_shuffle(cls, ds, batch_size: int, feats, labels):
     dl = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)
     last_X, last_Y = jnp.array([]), jnp.array([])
+    assert len(dl) == len(feats) // batch_size + 1
     for _ in range(2):
         X_list, Y_list = [], []
         for batch in dl:
@@ -65,6 +68,7 @@ def test_shuffle(cls, ds, batch_size: int, feats, labels):
 # %% ../nbs/tests.ipynb 7
 def test_shuffle_drop_last(cls, ds, batch_size: int, feats, labels):
     dl = cls(ds, batch_size=batch_size, shuffle=True, drop_last=True)
+    assert len(dl) == len(feats) // batch_size
     for _ in range(2):
         X_list, Y_list = [], []
         for batch in dl:

nbs/tests.ipynb

@@ -58,6 +58,7 @@
     "#| exporti\n",
     "def test_no_shuffle(cls, ds, batch_size: int, feats, labels):\n",
     "    dl = cls(ds, batch_size=batch_size, shuffle=False)\n",
+    "    assert len(dl) == len(feats) // batch_size + 1\n",
     "    for _ in range(2):\n",
     "        X_list, Y_list = [], []\n",
     "        for batch in dl:\n",
@@ -78,6 +79,7 @@
     "#| exporti\n",
     "def test_no_shuffle_drop_last(cls, ds, batch_size: int, feats, labels):\n",
     "    dl = cls(ds, batch_size=batch_size, shuffle=False, drop_last=True)\n",
+    "    assert len(dl) == len(feats) // batch_size\n",
     "    for _ in range(2):\n",
     "        X_list, Y_list = [], []\n",
     "        for batch in dl:\n",
@@ -100,6 +102,7 @@
     "def test_shuffle(cls, ds, batch_size: int, feats, labels):\n",
     "    dl = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)\n",
     "    last_X, last_Y = jnp.array([]), jnp.array([])\n",
+    "    assert len(dl) == len(feats) // batch_size + 1\n",
     "    for _ in range(2):\n",
     "        X_list, Y_list = [], []\n",
     "        for batch in dl:\n",
@@ -126,6 +129,7 @@
     "#| exporti\n",
     "def test_shuffle_drop_last(cls, ds, batch_size: int, feats, labels):\n",
     "    dl = cls(ds, batch_size=batch_size, shuffle=True, drop_last=True)\n",
+    "    assert len(dl) == len(feats) // batch_size\n",
     "    for _ in range(2):\n",
     "        X_list, Y_list = [], []\n",
     "        for batch in dl:\n",