Rename EagerArray to Tensor (#121)

Tensor seems a more appropriate name.

Author: gramalingam

onnxscript/autocast.py

@@ -1,5 +1,5 @@
 from onnx.defs import OpSchema
-from .eager_array import EagerArray
+from .tensor import Tensor
 import numpy as np
 
 
@@ -54,7 +54,7 @@ def cast_inputs(get_type_info, cast, opschema, *args):
 def dynamic_cast_inputs(opschema, *args):
     '''Used for autocast during eager-mode execution.'''
     def get_type_info(x):
-        return x.dtype if isinstance(x, EagerArray) else None
+        return x.dtype if isinstance(x, Tensor) else None
 
     def cast(x, typeinfo):
         if isinstance(x, (int, float)):
@@ -65,7 +65,7 @@ def cast(x, typeinfo):
                 dtype = np.int32
             else:  # isinstance(x, float):
                 dtype = np.float32
-            return EagerArray(np.array(x, dtype=dtype))
+            return Tensor(np.array(x, dtype=dtype))
         else:
             return x
 

onnxscript/eager_mode_evaluator.py

@@ -11,7 +11,7 @@
 from onnxruntime.capi.onnxruntime_pybind11_state import Fail, InvalidGraph, InvalidArgument
 from .utils import values_to_value_infos, proto2text
 from .irbuilder import select_ir_version
-from .eager_array import EagerArray
+from .tensor import Tensor
 
 
 class EagerModeError(RuntimeError):
@@ -61,7 +61,7 @@ def os_to_ort_value(v):
     '''
     Converts an onnxscript encoding of an ONNX value into the encoding used by ORT.
     '''
-    if isinstance(v, EagerArray):
+    if isinstance(v, Tensor):
         return v.value
     elif isinstance(v, list):
         return v
@@ -78,7 +78,7 @@ def ort_to_os_value(v):
     Converts an ORT encoding of an ONNX value into the encoding used by onnxscript.
     '''
     if isinstance(v, np.ndarray):
-        return EagerArray(v)
+        return Tensor(v)
     elif isinstance(v, list):
         return v
     elif v is None:

onnxscript/onnx_types.py

@@ -7,7 +7,7 @@
 import onnx.helper
 
 
-class Tensor:
+class TensorType:
     # Reference implementation placeholder
     # represents a generic ONNX tensor type
     def __init__(self, dtype=onnx.TensorProto.UNDEFINED, shape=None) -> None:
@@ -61,7 +61,7 @@ def mk_dim(dim):
             s = None
         else:
             s = [shape]
-        return Tensor(self.dtype, s)
+        return TensorType(self.dtype, s)
 
     def to_type_proto(self):
         return onnx.helper.make_tensor_type_proto(self.dtype, ())

onnxscript/eager_array.py renamed to onnxscript/tensor.py

@@ -7,30 +7,30 @@
 from onnx.mapping import NP_TYPE_TO_TENSOR_TYPE
 
 
-class EagerArray:
+class Tensor:
     """
-    Wraps arrays to intercept calls to operators and use onnxruntime
-    to process the output.
+    An implementation of ONNX Tensors, based on a wrapper around numpy arrays.
+    Serves to define overloaded ops with an ONNX/ONNXScript semantics.
     """
 
-    def __init__(self, tensor, opset=None):
-        if not isinstance(tensor, np.ndarray):
-            raise TypeError(f"Unexpected type {type(tensor)}. It must be a numpy array.")
-        self._tensor = tensor
+    def __init__(self, nparray, opset=None):
+        if not isinstance(nparray, np.ndarray):
+            raise TypeError(f"Unexpected type {type(nparray)}. It must be a numpy array.")
+        self._nparray = nparray
         from onnxscript.onnx_opset import default_opset
         self._opset = opset or default_opset
 
     @property
     def value(self):
-        return self._tensor
+        return self._nparray
 
     @property
     def shape(self):
-        return self._tensor.shape
+        return self._nparray.shape
 
     @property
     def dtype(self):
-        return self._tensor.dtype
+        return self._nparray.dtype
 
     @property
     def onnx_dtype(self):
@@ -52,7 +52,7 @@ def __getitem__(self, index):
         if isinstance(index, int):
             # case A[i]: indexing
             # promote integer input to tensor
-            i = EagerArray(np.array(index))
+            i = Tensor(np.array(index))
             # use Gather to perform indexing
             return op.Gather(self, i, axis=0)
         if not isinstance(index, (slice, tuple)):
@@ -78,13 +78,13 @@ def __getitem__(self, index):
             else:
                 raise TypeError(f"Unexpected type {type(s)}: slice or int expected.")
         indices = np.array(indices, dtype=np.int64).T
-        starts = EagerArray(indices[0])
-        ends = EagerArray(indices[1])
-        axis = EagerArray(indices[2])
-        steps = EagerArray(indices[3])
+        starts = Tensor(indices[0])
+        ends = Tensor(indices[1])
+        axis = Tensor(indices[2])
+        steps = Tensor(indices[3])
         result = op.Slice(self, starts, ends, axis, steps)
         if len(to_squeeze) > 0:
-            result = EagerArray(np.squeeze(result.value, axis=tuple(to_squeeze)))
+            result = Tensor(np.squeeze(result.value, axis=tuple(to_squeeze)))
         return result
 
     def __mod__(self, other):

onnxscript/utils.py

@@ -9,7 +9,7 @@
 import onnx
 from onnx import TensorProto, ValueInfoProto, ModelProto, FunctionProto
 from onnx.helper import make_tensor_type_proto, make_sequence_type_proto
-from .eager_array import EagerArray
+from .tensor import Tensor
 
 # print utility unavailable in ONNX 1.12 or earlier:
 try:
@@ -23,7 +23,7 @@ def value_to_type_proto(val):
     '''
     Return the ONNX type of a python-value.
     '''
-    if isinstance(val, (np.ndarray, EagerArray)):
+    if isinstance(val, (np.ndarray, Tensor)):
         elem_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[val.dtype]
         shape = val.shape
         return make_tensor_type_proto(elem_type, shape)

onnxscript/values.py

@@ -10,7 +10,7 @@
 from enum import IntFlag
 import numpy as np
 import onnx
-from .eager_array import EagerArray
+from .tensor import Tensor
 from .autocast import dynamic_cast_inputs
 
 
@@ -174,7 +174,7 @@ def __call__(self, *args, **kwargs):
         if len(args) == 0:
             # Operator Constant, it is usually called within a function.
             return self._libcall(**kwargs)
-        if isinstance(args[0], EagerArray):
+        if isinstance(args[0], Tensor):
             return self._libcall(*args, **kwargs)
         return self._usercall(*args, **kwargs)
 
@@ -183,21 +183,21 @@ def _usercall(self, *args, **kwargs):
         new_args = []
         for i, a in enumerate(args):
             if isinstance(a, np.ndarray):
-                new_args.append(EagerArray(a))
+                new_args.append(Tensor(a))
             elif isinstance(a, bool):
-                new_args.append(EagerArray(np.array(a)))
+                new_args.append(Tensor(np.array(a)))
             else:
                 raise TypeError(
                     f"Unexpected input type {type(a)} for an input {i}.")
         res = self.function(*new_args, **kwargs)
         if isinstance(res, np.ndarray):
             return res
-        if isinstance(res, EagerArray):
+        if isinstance(res, Tensor):
             return res.value
         if isinstance(res, (list, tuple)):
             unwrapped = []
             for i, r in enumerate(res):
-                if isinstance(r, EagerArray):
+                if isinstance(r, Tensor):
                     unwrapped.append(r.value)
                 else:
                     raise TypeError(
@@ -216,23 +216,23 @@ def _libcall(self, *args, **kwargs):
         """
         new_args = []
         for i, a in enumerate(args):
-            if isinstance(a, EagerArray):
+            if isinstance(a, Tensor):
                 new_args.append(a)
             elif isinstance(a, bool):
                 # TODO: default values for function parameters
                 # are not properly handled yet. This section
                 # should disappear.
-                new_args.append(EagerArray(np.array(a)))
+                new_args.append(Tensor(np.array(a)))
             else:
                 raise TypeError(
                     f"Unexpected input type {type(a)} for an input {i}.")
         res = self.function(*new_args, **kwargs)
-        if isinstance(res, EagerArray):
+        if isinstance(res, Tensor):
             return res
         if isinstance(res, tuple):
             unwrapped = []
             for i, r in enumerate(res):
-                if isinstance(r, EagerArray):
+                if isinstance(r, Tensor):
                     unwrapped.append(r)
                 else:
                     raise TypeError(