doc

Author: xadupre

onnx_diagnostic/export/onnx_plug.py

@@ -143,7 +143,7 @@ def forward(self, x):
         print(pretty_onnx(onx))
 
     This shows how to define multiple versions depending on the device,
-    the type or the targetted onnx opset.
+    the type or the targeted onnx opset.
 
     .. code-block:: python
 

onnx_diagnostic/helpers/onnx_helper.py

@@ -17,7 +17,6 @@
     Union,
 )
 import numpy as np
-import numpy.typing as npt
 import onnx
 import onnx.helper as oh
 import onnx.numpy_helper as onh
@@ -33,6 +32,8 @@
     load as onnx_load,
 )
 
+TensorLike = Union[np.ndarray, "torch.Tensor"]  # noqa: F821
+
 
 def _make_stat(init: TensorProto) -> Dict[str, float]:
     """
@@ -490,7 +491,7 @@ def convert_endian(tensor: TensorProto) -> None:
     tensor.raw_data = np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap().tobytes()
 
 
-def from_array_ml_dtypes(arr: npt.ArrayLike, name: Optional[str] = None) -> TensorProto:
+def from_array_ml_dtypes(arr: TensorLike, name: Optional[str] = None) -> TensorProto:
     """
     Converts a numpy array to a tensor def assuming the dtype
     is defined in ml_dtypes.
@@ -536,7 +537,7 @@ def from_array_ml_dtypes(arr: npt.ArrayLike, name: Optional[str] = None) -> Tens
 }
 
 
-def from_array_extended(tensor: npt.ArrayLike, name: Optional[str] = None) -> TensorProto:
+def from_array_extended(tensor: TensorLike, name: Optional[str] = None) -> TensorProto:
     """
     Converts an array into a :class:`onnx.TensorProto`.
 
@@ -603,7 +604,7 @@ def from_array_extended(tensor: npt.ArrayLike, name: Optional[str] = None) -> Te
     return t
 
 
-def to_array_extended(proto: TensorProto) -> npt.ArrayLike:
+def to_array_extended(proto: TensorProto) -> TensorLike:
     """Converts :class:`onnx.TensorProto` into a numpy array."""
     arr = onh.to_array(proto)
     if proto.data_type >= onnx.TensorProto.BFLOAT16:

onnx_diagnostic/helpers/ort_session.py

@@ -1,7 +1,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import onnx
 import numpy as np
-import numpy.typing as npt
 import torch
 from torch._C import _from_dlpack
 import onnxruntime
@@ -16,6 +15,7 @@
 
 
 DEVICES = {-1: ORTC.OrtDevice(ORTC.OrtDevice.cpu(), ORTC.OrtDevice.default_memory(), 0)}
+TensorLike = Union[np.ndarray, torch.Tensor]
 
 
 class _InferenceSession:
@@ -243,16 +243,16 @@ def __init__(
         )
 
     def run(
-        self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
-    ) -> List[Optional[npt.ArrayLike]]:
+        self, output_names: Optional[List[str]], feeds: Dict[str, TensorLike]
+    ) -> List[Optional[TensorLike]]:
         """Calls :meth:`onnxruntime.InferenceSession.run`."""
         # sess.run does not support blfoat16
         # res = self.sess.run(output_names, feeds)
         return self._post_process_inplace(list(self.run_dlpack(output_names, feeds)))
 
     def run_dlpack(
-        self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
-    ) -> Tuple[Optional[npt.ArrayLike], ...]:
+        self, output_names: Optional[List[str]], feeds: Dict[str, TensorLike]
+    ) -> Tuple[Optional[TensorLike], ...]:
         """
         Same as :meth:`onnxruntime.InferenceSession.run` except that
         feeds is a dictionary of :class:`np.ndarray`.
@@ -289,13 +289,13 @@ def run_dlpack(
     def _ortvalues_to_numpy_tensor(
         self,
         ortvalues: Union[List[ORTC.OrtValue], ORTC.OrtValueVector],
-    ) -> Tuple[Optional[npt.ArrayLike], ...]:
+    ) -> Tuple[Optional[TensorLike], ...]:
         if len(ortvalues) == 0:
             return tuple()
 
         if self.nvtx:
             self.torch.cuda.nvtx.range_push("_ortvalues_to_numpy_tensor")
-        res: List[Optional[npt.ArrayLike]] = []  # noqa: F823
+        res: List[Optional[TensorLike]] = []  # noqa: F823
         for i in range(len(ortvalues)):
             if not ortvalues[i].has_value():
                 res.append(None)
@@ -556,7 +556,7 @@ def investigate_onnxruntime_issue(
         Union[str, Callable[[onnx.ModelProto], onnxruntime.InferenceSession]]
     ] = None,
     # if model needs to be run.
-    feeds: Optional[Union[Dict[str, torch.Tensor], Dict[str, npt.ArrayLike]]] = None,
+    feeds: Optional[Union[Dict[str, torch.Tensor], Dict[str, TensorLike]]] = None,
     verbose: int = 0,
     dump_filename: Optional[str] = None,
     infer_shapes: bool = True,

onnx_diagnostic/helpers/torch_helper.py

@@ -1107,7 +1107,7 @@ def study_discrepancies(
 
 def int_device_to_torch_device(device_id: int) -> torch.device:
     """
-    Converts a device defined as an integer (coming from :meth:`torch.get_device`)
+    Converts a device defined as an integer (coming from :meth:`torch.Tensor.get_device`)
     into a ``torch.device``.
     """
     if device_id < 0: