test array write

huggingface · Oct 15, 2024 · c39c4bc · c39c4bc
1 parent abbb59a
commit c39c4bc
Showing 1 changed file with 11 additions and 11 deletions.
diff --git a/tests/features/test_array_xd.py b/tests/features/test_array_xd.py
@@ -175,20 +175,20 @@ def get_dict_examples(self, shape_1, shape_2):
     def _check_getitem_output_type(self, dataset, shape_1, shape_2, first_matrix):
         matrix_column = dataset["matrix"]
         self.assertIsInstance(matrix_column, list)
-        self.assertIsInstance(matrix_column[0], list)
-        self.assertIsInstance(matrix_column[0][0], list)
+        self.assertIsInstance(matrix_column[0], np.ndarray)
+        self.assertIsInstance(matrix_column[0][0], np.ndarray)
         self.assertTupleEqual(np.array(matrix_column).shape, (2, *shape_2))
 
         matrix_field_of_first_example = dataset[0]["matrix"]
-        self.assertIsInstance(matrix_field_of_first_example, list)
-        self.assertIsInstance(matrix_field_of_first_example, list)
+        self.assertIsInstance(matrix_field_of_first_example, np.ndarray)
+        self.assertIsInstance(matrix_field_of_first_example[0], np.ndarray)
         self.assertEqual(np.array(matrix_field_of_first_example).shape, shape_2)
         np.testing.assert_array_equal(np.array(matrix_field_of_first_example), np.array(first_matrix))
 
         matrix_field_of_first_two_examples = dataset[:2]["matrix"]
         self.assertIsInstance(matrix_field_of_first_two_examples, list)
-        self.assertIsInstance(matrix_field_of_first_two_examples[0], list)
-        self.assertIsInstance(matrix_field_of_first_two_examples[0][0], list)
+        self.assertIsInstance(matrix_field_of_first_two_examples[0], np.ndarray)
+        self.assertIsInstance(matrix_field_of_first_two_examples[0][0], np.ndarray)
         self.assertTupleEqual(np.array(matrix_field_of_first_two_examples).shape, (2, *shape_2))
 
         with dataset.formatted_as("numpy"):
@@ -268,7 +268,7 @@ def test_to_pylist(self):
         pylist = arr_xd.to_pylist()
 
         for first_dim, single_arr in zip(first_dim_list, pylist):
-            self.assertIsInstance(single_arr, list)
+            self.assertIsInstance(single_arr, np.ndarray)
             self.assertTupleEqual(np.array(single_arr).shape, (first_dim, *fixed_shape))
 
     def test_to_numpy(self):
@@ -311,8 +311,8 @@ def test_iter_dataset(self):
 
         for first_dim, ds_row in zip(first_dim_list, dataset):
             single_arr = ds_row["image"]
-            self.assertIsInstance(single_arr, list)
-            self.assertTupleEqual(np.array(single_arr).shape, (first_dim, *fixed_shape))
+            self.assertIsInstance(single_arr, np.ndarray)
+            self.assertTupleEqual(single_arr.shape, (first_dim, *fixed_shape))
 
     def test_to_pandas(self):
         fixed_shape = (2, 2)
@@ -353,8 +353,8 @@ def test_map_dataset(self):
         # check also if above function resulted with 2x bigger first dim
         for first_dim, ds_row in zip(first_dim_list, dataset):
             single_arr = ds_row["image"]
-            self.assertIsInstance(single_arr, list)
-            self.assertTupleEqual(np.array(single_arr).shape, (first_dim * 2, *fixed_shape))
+            self.assertIsInstance(single_arr, np.ndarray)
+            self.assertTupleEqual(single_arr.shape, (first_dim * 2, *fixed_shape))
 
 
 @pytest.mark.parametrize("dtype, dummy_value", [("int32", 1), ("bool", True), ("float64", 1)])