Update imports to use numpy

Author: BirkhoffG

integration_tests/tf_test.py

@@ -1,5 +1,5 @@
 import jax_dataloader as jdl
-import jax
+import numpy as np
 import tensorflow_datasets as tfds
 import tensorflow as tf
 
@@ -9,13 +9,13 @@ def test_jax():
     dl = jdl.DataLoader(ds, 'jax', batch_size=2)
     for x, y in dl:
         z = x + y
-        assert isinstance(z, jax.Array)
+        assert isinstance(z, np.ndarray)
 
 
 def test_tf():
     ds = tf.data.Dataset.from_tensor_slices((tf.ones((10, 3)), tf.ones((10, 3))))
     dl = jdl.DataLoader(ds, 'tensorflow', batch_size=2)
     for x, y in dl: 
         z = x + y
-        assert isinstance(z, jax.Array)
+        assert isinstance(z, np.ndarray)