refactoring

Author: xadupre

_doc/api/api.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.api
+===================
+
+.. automodule:: onnx_diagnostic.api
+    :members:
+    :no-undoc-members:

_doc/api/index.rst

@@ -19,6 +19,7 @@ API of onnx_diagnostic
     :maxdepth: 1
     :caption: modules
 
+    api
     ext_test_case
 
 .. automodule:: onnx_diagnostic

_doc/api/reference/torch_ops/index.rst

@@ -24,16 +24,22 @@ OpRun
 .. autoclass:: onnx_diagnostic.reference.torch_ops.OpRun
     :members:
 
+OpRunTensor
++++++++++++
+
+.. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunTensor
+    :members:
+
 OpRunValue
 ++++++++++
 
 .. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunValue
     :members:
 
-OpRunValueSequence
-++++++++++++++++++
+OpRunSequence
++++++++++++++
 
-.. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunValueSequence
+.. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunSequence
     :members:
 
 OpRunFunction

onnx_diagnostic/api.py

@@ -0,0 +1,2 @@
+class TensorLike:
+    """Mocks a tensor."""

onnx_diagnostic/reference/torch_evaluator.py

@@ -264,7 +264,7 @@ def run(
             r = self.runtime_info[k]
             if not r.has_value:
                 r.set_value(
-                    torch_ops.OpRunValue(
+                    torch_ops.OpRunTensor(
                         v.to(self.CUDA) if not r.is_shape and self.on_cuda else v,
                         is_constant=True,
                         may_cpu=len(v.shape) == 1 and v.numel() < 8 and v.dtype == torch.int64,
@@ -277,7 +277,7 @@ def run(
         for k, v in feeds.items():
             r = self.runtime_info[k]
             r.set_value(
-                torch_ops.OpRunValue(
+                torch_ops.OpRunTensor(
                     v.to(self.CUDA) if not r.is_shape and self.on_cuda else v,
                     is_constant=False,
                     may_cpu=len(v.shape) == 1 and v.numel() < 8 and v.dtype == torch.int64,
@@ -352,15 +352,15 @@ def run(
 
     def run_with_values(
         self,
-        *args: Optional[torch_ops.OpRunValue],
+        *args: Optional[torch_ops.OpRunTensor],
         context: Optional[Dict[str, RuntimeValue]] = None,
     ) -> Union[torch_ops.OpRunValue, Tuple[torch_ops.OpRunValue, ...]]:
         """
         Runs the ONNX model.
 
         :param args: inputs
         :param context: local context for the execution of subgraphs
-        :return: output OpRunValue
+        :return: output OpRunTensor
         """
         assert all(
             isinstance(a, torch_ops.OpRunValue) for a in args
@@ -373,7 +373,7 @@ def run_with_values(
             r = self.runtime_info[k]
             if not r.has_value:
                 r.set_value(
-                    torch_ops.OpRunValue(
+                    torch_ops.OpRunTensor(
                         v.to(self.CUDA) if r.is_shape is False and self.on_cuda else v,
                         is_constant=True,
                         may_cpu=len(v.shape) == 1 and v.numel() < 8 and v.dtype == torch.int64,
@@ -384,7 +384,7 @@ def run_with_values(
         for k, v in zip(self.input_names, args):
             r = self.runtime_info[k]
             r.set_value(
-                torch_ops.OpRunValue(None) if v is None else v.__class__(v.tensor_or_sequence)
+                torch_ops.OpRunTensor(None) if v is None else v.__class__(v.tensor_or_sequence)
             )
 
         # node execution
@@ -406,7 +406,7 @@ def run_with_values(
                 res = kernel.run(*inputs)
                 if isinstance(res, tuple):
                     # outputs
-                    assert all(isinstance(o, torch_ops.OpRunValue) for o in res), (
+                    assert all(isinstance(o, torch_ops.OpRunTensor) for o in res), (
                         f"Unexpected output type {[type(o) for o in res]} "
                         f"for kernel {type(kernel)}."
                     )
@@ -425,7 +425,7 @@ def run_with_values(
         assert all(
             self.runtime_info[o].value is not None for o in outputs
         ), "Not implemented yet when one output is None."
-        res2 = [torch_ops.OpRunValue(self.runtime_info[o].value.tensor) for o in outputs]  # type: ignore[assignment, union-attr]
+        res2 = [self.runtime_info[o].value.copy() for o in outputs]  # type: ignore[assignment, union-attr]
 
         # clean previous execution
         for k in self.input_names:

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -1,4 +1,4 @@
-from ._op_run import OpRun, OpRunFunction, OpRunValue, OpRunValueSequence
+from ._op_run import OpRun, OpRunFunction, OpRunSequence, OpRunTensor, OpRunValue
 from .access_ops import Gather_1, Slice_13
 from .binary_ops import (
     And_1,

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -1,11 +1,23 @@
 from typing import Any, List, Optional, Union, Tuple
 import onnx
 import torch
+from ...api import TensorLike
 from ...helpers import string_type
 from ...helpers.torch_helper import to_tensor
 
 
-class OpRunValue:
+class OpRunValue(TensorLike):
+    """Defines a value for the runtime, a tensor or a sequence."""
+
+    __slots__ = ("cached", "is_constant", "sequence", "tensor")
+
+    @classmethod
+    def is_sequence(cls) -> bool:
+        "Tells if it is sequence."
+        raise NotImplementedError("is_sequence must be overwritten.")
+
+
+class OpRunTensor(OpRunValue):
     """
     Wrapper around a tensor.
 
@@ -15,10 +27,9 @@ class OpRunValue:
         more appropriate
     """
 
-    __slots__ = ("cached", "is_constant", "tensor")
-
     def __init__(self, tensor, is_constant: bool = False, may_cpu: bool = False):
         assert isinstance(tensor, torch.Tensor), f"Unexpected type {type(tensor)}"
+        assert tensor is None or tensor.numel() > 1 or tensor.item() != -666666
         self.tensor = (
             tensor.cpu()
             if may_cpu
@@ -36,9 +47,9 @@ def is_sequence(cls) -> bool:
         "Tells if it is sequence."
         return False
 
-    def to(self, to: Any) -> "OpRunValue":
+    def to(self, to: Any) -> "OpRunTensor":
         "Changes the device."
-        return OpRunValue(self.tensor.to(to))
+        return OpRunTensor(self.tensor.to(to))
 
     def string_type(self) -> str:
         "Returns information about the value as a string."
@@ -96,17 +107,29 @@ def as_tuple_int(self) -> Tuple[int, ...]:
             return self.cached
         return self._tensor_as_tuple_int()
 
+    def copy(self) -> "OpRunTensor":
+        "Shallow copy."
+        return self.__class__(self.tensor)
 
-class OpRunValueSequence(OpRunValue):
-    """Defines a sequence."""
 
-    __slots__ = ("cached", "is_constant", "sequence", "tensor")
+class OpRunSequence(OpRunValue):
+    """Defines a sequence."""
 
     def __init__(
         self, sequence: Optional[List[torch.Tensor]] = None, dtype: torch.dtype = torch.float32
     ):
-        super().__init__(torch.empty((), dtype=dtype), False, False)
+        self.tensor = torch.tensor(-666666, dtype=dtype)
+        self.is_shape = False
         self.sequence = sequence or []
+        self.cached: Optional[Tuple[int, ...]] = None
+        assert all(
+            isinstance(s, torch.Tensor) for s in self.sequence
+        ), f"Unexpected type in sequence {[type(s) for s in self.sequence]}"
+
+    @property
+    def dtype(self):
+        "dtype (torch dtype)"
+        return self.tensor.dtype
 
     @property
     def tensor_or_sequence(self) -> Union[torch.Tensor, List[torch.Tensor]]:
@@ -119,18 +142,23 @@ def is_sequence(cls) -> bool:
         return True
 
     def insert_at(
-        self, tensor: OpRunValue, position: Optional[OpRunValue] = None
-    ) -> "OpRunValueSequence":
+        self, tensor: torch.Tensor, position: Optional[OpRunTensor] = None
+    ) -> "OpRunSequence":
         "Inserts a value at a given position."
-        new_seq = OpRunValueSequence()
+        assert isinstance(tensor, OpRunTensor), f"Unexpected type {type(tensor)} for tensor"
+        new_seq = OpRunSequence()
         seq = self.sequence.copy()
         new_seq.sequence = seq
         if position is None:
-            seq.append(tensor)
+            seq.append(tensor.tensor)
         else:
-            seq.insert(int(position.tensor.item()), tensor)
+            seq.insert(int(position.tensor.item()), tensor.tensor)
         return new_seq
 
+    def copy(self) -> "OpRunSequence":
+        "Shallow copy."
+        return self.__class__(self.sequence, dtype=self.dtype)
+
 
 class OpRun:
     """

onnx_diagnostic/reference/torch_ops/access_ops.py

@@ -1,7 +1,7 @@
 from typing import Optional
 import onnx
 import torch
-from . import OpRun, OpRunValue
+from . import OpRun, OpRunTensor
 
 
 class Gather_1(OpRun):
@@ -18,20 +18,20 @@ def run(self, x, indices):
             return torch.empty((0,), dtype=x.tensor.dtype, device=x.tensor.device)
         ind = [slice(0, s) for s in x.shape]
         ind[self.axis] = indices.tensor
-        return OpRunValue(x.tensor[tuple(ind)])
+        return OpRunTensor(x.tensor[tuple(ind)])
 
 
 class Slice_13(OpRun):
     "Slice"
 
     def run(
         self,
-        data: OpRunValue,
-        starts: OpRunValue,
-        ends: OpRunValue,
-        axes: Optional[OpRunValue] = None,
-        steps: Optional[OpRunValue] = None,
-    ) -> OpRunValue:
+        data: OpRunTensor,
+        starts: OpRunTensor,
+        ends: OpRunTensor,
+        axes: Optional[OpRunTensor] = None,
+        steps: Optional[OpRunTensor] = None,
+    ) -> OpRunTensor:
         if axes is None:
             if steps is None:
                 slices = [slice(s, e) for s, e in zip(starts.tensor, ends.tensor)]
@@ -48,4 +48,4 @@ def run(
                 slices = [slice(0, a) for a in data.shape]
                 for s, e, a, d in zip(starts.tensor, ends.tensor, axes.tensor, steps.tensor):
                     slices[a] = slice(s, e, d)
-        return OpRunValue(data.tensor[tuple(slices)])
+        return OpRunTensor(data.tensor[tuple(slices)])