study

Author: xadupre

_unittests/ut_helpers/test_torch_helper.py

@@ -28,7 +28,7 @@
 )
 from onnx_diagnostic.helpers.mini_onnx_builder import create_input_tensors_from_onnx_model
 from onnx_diagnostic.helpers.onnx_helper import from_array_extended, to_array_extended
-from onnx_diagnostic.helpers.torch_helper import to_tensor
+from onnx_diagnostic.helpers.torch_helper import to_tensor, study_discrepancies
 
 TFLOAT = onnx.TensorProto.FLOAT
 
@@ -425,6 +425,12 @@ def test_get_weight_type(self):
         dt = get_weight_type(model)
         self.assertEqual(torch.float32, dt)
 
+    def test_study_discrepancies(self):
+        t1 = torch.rand((3, 4))
+        t2 = torch.rand((3, 4))
+        ax = study_discrepancies(t1, t2)
+        self.assertEqual(ax.shape, ((3, 2)))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/helpers/torch_helper.py

@@ -1024,7 +1024,7 @@ def study_discrepancies(
     figsize: Optional[Tuple[int, int]] = (15, 15),
     title: Optional[str] = None,
     name: Optional[str] = None,
-) -> "Axes":  # noqa: F821
+) -> "matplotlib.axes.Axes":  # noqa: F821
     """
     Computes different metrics for the discrepancies.
     Returns graphs.