tensor.py

"""Defines Tensor object."""
import numpy as np


class Tensor(object):
  """Symbolic representation of a multi-dimensional array.

  Tensor is always generated by an Operation. Tensor can be uniquely identified
  by the parent Op (`self.op`), and its index in the output list of parent Op (
  `tensor_index`). For example, an "Add" Op generates an output Tensor with
  `tensor_index` 0, while an "Unpack" Op genrates output Tensor`s with
  `tensor_index` 0, 1, 2 ...

  Tensors are considered "immutable" in the sense that their value is solely
  determined by their parent Op, and therefore cannot be modified.

  Each tensor has a `shape` attribute that indicates the shape of the multi
  dimensional array at graph construction time. See `TensorShape` for details.
  """

  def __init__(self, op, tensor_index, shape):
    """Constructor.

    Args:
      op (Operation): the parent Op.
      tensor_index (int): tensor's index in the output list of `op`.
      shape (TensorShape): shape of the `Tensor`.
    """
    from .operation import Operation
    if not isinstance(op, Operation):
      raise ValueError(f"op must be an Operation, got {type(op)}")

    self._op = op
    self._tensor_index = tensor_index
    self._shape = shape

  @property
  def op(self):
    """Parent Op."""
    return self._op

  @property
  def tensor_index(self):
    """Index of tensor in the output list of its parent Op."""
    return self._tensor_index

  @property
  def shape(self):
    """The `TensorShape`."""
    return self._shape

  @property
  def name(self):
    """String that uniquely identify this tensor."""
    return f"{self.op.name}:{self.tensor_index}"

  def __repr__(self):
    repstr = (
        f"<Tensor '{self.op.name}:{self.tensor_index}', "
        f"shape={self.shape.raw_shape}>"
    )
    return repstr

  def _convert_arithmetic_operand(self, other):
    from .generic_ops import Const
    if not isinstance(other, Tensor):
      try:
        other = Const(value=np.asarray(other)).output(0)
      except Exception:
        raise TypeError('other must be a Tensor or convertable to numpy array.')
    return other

  def __add__(self, other):
    from .math_ops import Add
    other = self._convert_arithmetic_operand(other)
    return Add(input_list=[self, other]).output(0)

  def __radd__(self, other):
    from .math_ops import Add
    other = self._convert_arithmetic_operand(other)
    return Add(input_list=[self, other]).output(0)

  def __mul__(self, other):
    from .math_ops import Mul
    other = self._convert_arithmetic_operand(other)
    return Mul(input_list=[self, other]).output(0)

  def __rmul__(self, other):
    from .math_ops import Mul
    other = self._convert_arithmetic_operand(other)
    return Mul(input_list=[self, other]).output(0)

  def __sub__(self, other):
    from .math_ops import Sub
    other = self._convert_arithmetic_operand(other)
    return Sub(input_list=[self, other]).output(0)

  def __rsub__(self, other):
    from .math_ops import Sub
    other = self._convert_arithmetic_operand(other)
    return Sub(input_list=[other, self]).output(0)

  def __div__(self, other):
    from .math_ops import RealDiv
    other = self._convert_arithmetic_operand(other)
    return RealDiv(input_list=[self, other]).output(0)

  def __rdiv__(self, other):
    from .math_ops import RealDiv
    other = self._convert_arithmetic_operand(other)
    return RealDiv(input_list=[other, self]).output(0)

  def __pos__(self):
    return self

  def __neg__(self):
    from .math_ops import Neg
    return Neg(input_list=[self]).output(0)

  def __getitem__(self, slice_specs):
    from .array_ops import Reshape, StridedSlice, Unpack
    from .generic_ops import Const

    ndims = None
    orig_shape = None
    if self.shape.level > 0:
      ndims = self.shape.ndims
      orig_shape = list(self.shape.raw_shape)
      if self.shape.level == 1:
        # replace Nones in `orig_shape` with dynamic axis size
        shape = self.op.get_shape_tensor(tensor_index=self.tensor_index)
        unpack_op = Unpack(input_list=[shape], num=ndims, axis=0)
        for i in np.arange(ndims):
          s = unpack_op.output(i)
          if orig_shape[i] is None:
            orig_shape[i] = s

    new_shape = []

    if not isinstance(slice_specs, tuple):
      slice_specs = (slice_specs,)

    begin = []
    end = []
    strides = []
    ellipsis_index = None
    new_axis_count = 0

    for index, ss in enumerate(slice_specs):
      if isinstance(ss, int):
        begin.append(ss)
        end.append(ss + 1)
        strides.append(1)
        new_shape.append(orig_shape[index - new_axis_count])

      elif isinstance(ss, slice):
        begin.append(0 if ss.start is None else ss.start)
        end.append(
            orig_shape[index - new_axis_count] if ss.stop is None else ss.stop,
        )
        strides.append(1 if ss.step is None else ss.step)
        new_shape.append(orig_shape[index - new_axis_count])

      elif ss is None:
        begin.append(0)
        end.append(1)
        strides.append(1)
        new_shape.append(1)
        new_axis_count += 1
      elif ss is Ellipsis:
        raise NotImplementedError(
            "Ellipsis is currently not supported for slicing.",
        )

    pad_size = len(orig_shape) - (index + 1 - new_axis_count)
    begin = begin + [0] * pad_size
    end = end + orig_shape[index + 1 - new_axis_count:len(orig_shape)]
    strides = strides + [1] * pad_size
    new_shape = new_shape + orig_shape[
        index + 1 -
        new_axis_count:len(orig_shape)
    ]

    new_shape = _build_vector_from_mixed(new_shape)

    tensor = Reshape(input_list=[self, new_shape]).output(0)

    begin = Const(value=np.asarray(begin)).output(0)
    end = _build_vector_from_mixed(end)
    strides = Const(value=np.asarray(strides)).output(0)

    tensor = StridedSlice(input_list=[tensor, begin, end, strides]).output(0)
    return tensor

  def eval(self):
    """Return the value of the Tensor."""
    self.op.run()
    runtime = self.op.graph.runtime
    tensor_value = runtime._values[self.op.id][self.tensor_index]
    return tensor_value


class Placeholder(Tensor):
  """A `Placeholder` is a special type of Tensor such that its shape must be
  determined at graph construction time, and its value is not known until the
  runtime (by `set_value`).
  """

  def set_value(self, value):
    """Set the value to be used by the placeholder at runtime.

    Args:
      value (numpy array like): value convertible to numpy array.
    """
    runtime = self.op.graph.runtime
    runtime._placeholder_values[self.op.id] = value


def _build_vector_from_mixed(mixed):
  """Create a 1-D tensor for a list of Tensors or numeric values.

  Args:
    mixed (List): a list of Tensor or numeric values.

  Returns:
    vector (Tensor): a Const or Pack Tensor.
  """

  from .array_ops import Pack
  from .generic_ops import Const

  if not any([isinstance(i, Tensor) for i in mixed]):
    vector = Const(value=np.asarray(mixed)).output(0)
  else:
    tensorized = [
        i if isinstance(i, Tensor) else Const(value=np.asarray(i)).output(0)
        for i in mixed
    ]
    vector = Pack(input_list=tensorized, axis=0).output(0)
  return vector