prep.

Author: sayakpaul

tests/quantization/bnb/test_4bit.py

@@ -29,6 +29,7 @@
     is_transformers_available,
     load_pt,
     numpy_cosine_similarity_distance,
+    print_tensor_test,
     require_accelerate,
     require_bitsandbytes_version_greater,
     require_torch,
@@ -432,7 +433,7 @@ def test_quality(self):
         expected_slice = np.array([0.1123, 0.1296, 0.1609, 0.1042, 0.1230, 0.1274, 0.0928, 0.1165, 0.1216])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
-        print(f"{max_diff=}")
+        print_tensor_test(out_slice)
         self.assertTrue(max_diff < 1e-2)
 
     def test_generate_quality_dequantize(self):
@@ -448,6 +449,7 @@ def test_generate_quality_dequantize(self):
         ).images
 
         out_slice = output[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.1216, 0.1387, 0.1584, 0.1152, 0.1318, 0.1282, 0.1062, 0.1226, 0.1228])
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
         self.assertTrue(max_diff < 1e-3)
@@ -521,6 +523,7 @@ def test_quality(self):
         ).images
 
         out_slice = output[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.0583, 0.0586, 0.0632, 0.0815, 0.0813, 0.0947, 0.1040, 0.1145, 0.1265])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)

tests/quantization/bnb/test_mixed_int8.py

@@ -26,6 +26,7 @@
     is_transformers_available,
     load_pt,
     numpy_cosine_similarity_distance,
+    print_tensor_test,
     require_accelerate,
     require_bitsandbytes_version_greater,
     require_torch,
@@ -369,6 +370,7 @@ def test_quality(self):
             output_type="np",
         ).images
         out_slice = output[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.0149, 0.0322, 0.0073, 0.0134, 0.0332, 0.011, 0.002, 0.0232, 0.0193])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
@@ -418,6 +420,7 @@ def test_generate_quality_dequantize(self):
         ).images
 
         out_slice = output[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.0266, 0.0264, 0.0271, 0.0110, 0.0310, 0.0098, 0.0078, 0.0256, 0.0208])
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
         self.assertTrue(max_diff < 1e-2)
@@ -468,6 +471,7 @@ def test_quality(self):
             output_type="np",
         ).images
         out_slice = output[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.0574, 0.0554, 0.0581, 0.0686, 0.0676, 0.0759, 0.0757, 0.0803, 0.0930])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)