message: Define Tensor as a generic class
diff: @@ -24,6 +24,7 @@ import sys
  import threading
  import types
  +from typing import Generic, TypeVar
  import numpy as np
  import six
  from six.moves import map # pylint: disable=redefined-builtin
  @@ -253,9 +254,24 @@ def disable_tensor_equality():
  Tensor._USE_EQUALITY = False # pylint: disable=protected-access
  +DataType = TypeVar("DataType",
  + dtypes.Float16, dtypes.Float32, dtypes.Float64,
  + dtypes.BFloat16, dtypes.Complex64, dtypes.Complex128,
  + dtypes.Int8, dtypes.UInt8, dtypes.UInt16,
  + dtypes.UInt32, dtypes.UInt64, dtypes.Int16,
  + dtypes.Int32, dtypes.Int64, dtypes.Bool,
  + dtypes.String, dtypes.QInt8, dtypes.QUInt8,
  + dtypes.QInt16, dtypes.QUInt16, dtypes.QInt32,
  + dtypes.Resource, dtypes.Variant)
  +
  +if sys.version_info[0] >= 3:
  + GenericTensor = Generic[DataType]
  +else:
  + GenericTensor = object
  +
  # TODO(mdan): This object should subclass Symbol, not just Tensor.
  @tf_export("Tensor")
  -class Tensor(internal.NativeObject, core_tf_types.Tensor):
  +class Tensor(internal.NativeObject, core_tf_types.Tensor, GenericTensor):
  """A tensor is a multidimensional array of elements represented by a
  `tf.Tensor` object. All elements are of a single known data type.
  