text: import numpy as np
  import tensorflow as tf
  
  from t3f.tensor_train_base import TensorTrainBase
  from t3f.tensor_train import TensorTrain
  from t3f import shapes
  
  
  class TensorTrainBatch(TensorTrainBase):
    """Represents a batch of Tensor Train objects (TT-tensors or TT-matrices).
  
    t3f represents a Tensor Train object as a tuple of TT-cores.
    """
  
    def __init__(self, tt_cores, shape=None, tt_ranks=None, batch_size=None,
                 convert_to_tensors=True, name="TensorTrainBatch"):
      """Creates a `TensorTrainBatch`.
  
      Args:
        tt_cores: A tuple of 4d or 5d tensor-like objects of shape
          `[batch_size, r_k-1, n_k, r_k]` or
          `[batch_size, r_k-1, n_k, m_k, r_k]`
          Tensor-like can be numpy array, tf.Tensor, of tf.Variable
        batch_size: number of elements in the batch. If None, tries to infer from
          the TT-cores (not always possible even if it should be, e.g. if ranks
          are unknown, than the whole shape of a core can be unknown).
        shape: Shape of the underlying tensor. If None, tries to infer from the
          TT-cores.
        tt_ranks: a TensorShape of length d+1 (d is the dimensionality of
          the underlying tensor). The first and the last ranks are assumed to
          equal to 1. If None, tries to infer the ranks from the cores.
        convert_to_tensors: bool, if True than convert each element of the
          tt_cores tuple into a tf.Tensor (e.g. to initialize from np.array)
        name: The name of ops.
  
      Returns:
        A `TensorTrainBatch`.
  
      Raises:
        ValueError if the provided TT-cores are not valid or inconsistent with
          the provided shape.
      """
      tt_cores = list(tt_cores)
      if convert_to_tensors:
        with tf.name_scope(name):
          for i in range(len(tt_cores)):
            name = "core%d" % i
            tt_cores[i] = tf.convert_to_tensor(tt_cores[i], name=name)
  
      if not _are_batch_tt_cores_valid(tt_cores, shape, tt_ranks, batch_size):
        raise ValueError('The tt_cores provided to TensorTrainBatch constructor '
                         'are not valid, have different dtypes, or are '
                         'inconsistent with the provided batch_size, shape, or '
                         'TT-ranks.')
  
      self._tt_cores = tuple(tt_cores)
      if batch_size is None:
        self._batch_size = tt_cores[0].shape.as_list()[0]
      else:
        self._batch_size = batch_size
      self._raw_shape = shapes.clean_raw_shape(shape)
      if self._raw_shape is None:
        self._raw_shape = _infer_batch_raw_shape(self._tt_cores)
      self._tt_ranks = None if tt_ranks is None else tf.TensorShape(tt_ranks)
      if self._tt_ranks is None:
        self._tt_ranks = _infer_batch_tt_ranks(self._tt_cores)
  
    def get_shape(self):
      """Get the `TensorShape` representing the shape of the dense tensor.
  
      The first dimension is the batch_size.
  
      Returns:
        A `TensorShape` object.
      """
      shape = TensorTrainBase.get_shape(self)
      return tf.TensorShape(np.hstack((self.batch_size, shape)))
  
    @property
    def tt_cores(self):
      """A tuple of TT-cores.
  
      Returns:
        A tuple of 4d or 5d tensors shape
          `[batch_size, r_k-1, n_k, r_k]`
        or
          `[batch_size, r_k-1, n_k, m_k, r_k]`
      """
      return self._tt_cores
  
    @property
    def batch_size(self):
      """The number of elements or None if not known."""
      return self._batch_size
  
    @property
    def left_tt_rank_dim(self):
      """The dimension of the left TT-rank in each TT-core."""
      return 1
  
    @property
    def right_tt_rank_dim(self):
      """The dimension of the right TT-rank in each TT-core."""
      if self.is_tt_matrix():
        # The dimensions of each TT-core are
        # [batch_idx, left_rank, n, m, right_rank]
        return 4
      else:
        # The dimensions of each TT-core are
        # [batch_idx, left_rank, n, right_rank]
        return 3
  
    def __str__(self):
      """A string describing the TensorTrainBatch, its TT-rank and shape."""
      shape = self.get_shape()
      tt_ranks = self.get_tt_ranks()
  
      if self.batch_size is None:
          batch_size_str = '(?)'
      else:
          batch_size_str = str(self.batch_size)
  
      if self.is_tt_matrix():
        raw_shape = self.get_raw_shape()
        type_str = 'TT-matrix variables' if self.is_variable() else 'TT-matrices'
  
        return "A %s element batch of %s of size %d x %d, underlying tensor " \
               "shape: %s x %s, TT-ranks: %s" % (batch_size_str, type_str,
                                                 shape[1], shape[2],
                                                 raw_shape[0], raw_shape[1],
                                                 tt_ranks)
      else:
        if self.is_variable():
          type_str = 'Tensor Train variables'
        else:
          type_str = 'Tensor Trains'
        return "A %s element batch of %s of shape %s, TT-ranks: %s" % \
               (batch_size_str, type_str, shape[1:], tt_ranks)
  
    @staticmethod
    def _do_collapse_dim(slice_spec):
      # Returns true if slice_spec is specified exactly and we want to collapse
      # the corresponding axis, i.e. return an object with less dims. To be used
      # in indexing functions.
      # If its a actual slice, nothing to collapse. Otherwise (a number or
      # a tf.Tensor) want to collapse.
      return not isinstance(slice_spec, slice)
  
    def _batch_dim_getitem(self, element_spec):
      """__getitem__ when provided only one (batch) index.
  
      Examples:
        a[1]
        a[1:3]
      """
  
      # This object index is specified exactly and we want to collapse the
      # batch_size axis, i.e. return a TensorTrain instead of a TensorTrainBatch.
      do_collapse_batch_dim = self._do_collapse_dim(element_spec)
  
      new_tt_cores = []
      for core_idx in range(self.ndims()):
        curr_core = self.tt_cores[core_idx]
        if self.is_tt_matrix():
          new_tt_cores.append(curr_core[element_spec, :, :, :, :])
        else:
          new_tt_cores.append(curr_core[element_spec, :, :, :])
      if do_collapse_batch_dim:
        # This index is specified exactly and we want to collapse the batch_size
        # axis, i.e. return a TensorTrain instead of a TensorTrainBatch.
        return TensorTrain(new_tt_cores, self.get_raw_shape(),
                           self.get_tt_ranks())
      else:
        batch_size = new_tt_cores[0].shape.as_list()[0]
        return TensorTrainBatch(new_tt_cores, self.get_raw_shape(),
                                self.get_tt_ranks(), batch_size)
  
    def _full_getitem(self, slice_spec):
      """__getitem__ when provided full index of length ndims + 1.
  
      Examples:
        a = t3f.random_tensor_batch((2, 3, 4), batch_size=5)
        a[:3, 1:2, 4, :]
      """
      if len(slice_spec) != self.ndims() + 1:
        raise ValueError('Expected %d indices, got %d' % (self.ndims() + 1,
                                                          len(slice_spec)))
      # This object index is specified exactly and we want to collapse the
      # batch_size axis, i.e. return a TensorTrain instead of a TensorTrainBatch.
      do_collapse_batch_dim = self._do_collapse_dim(slice_spec[0])
      remainder = None
      new_tt_cores = []
      for core_idx in range(self.ndims()):
        curr_core = self.tt_cores[core_idx]
        if self.is_tt_matrix():
          raise NotImplementedError
        else:
          sliced_core = curr_core[slice_spec[0], :, slice_spec[core_idx + 1], :]
          do_collapse_curr_dim = self._do_collapse_dim(slice_spec[core_idx + 1])
          if do_collapse_curr_dim:
            # This index is specified exactly and we want to collapse this axis.
            if remainder is None:
              remainder = sliced_core
            else:
              if do_collapse_batch_dim:
                remainder = tf.einsum('ab,bd->ad', remainder, sliced_core)
              else:
                remainder = tf.einsum('oab,obd->oad', remainder, sliced_core)
          else:
            if remainder is not None:
              # Add reminder from the previous collapsed cores to the current
              # core.
              if do_collapse_batch_dim:
                sliced_core = tf.einsum('ab,bid->aid', remainder, sliced_core)
              else:
                sliced_core = tf.einsum('oab,obid->oaid', remainder,
                                        sliced_core)
              remainder = None
            new_tt_cores.append(sliced_core)
  
      if remainder is not None:
        # The reminder obtained from collapsing the last cores.
        if do_collapse_batch_dim:
          new_tt_cores[-1] = tf.einsum('aib,bd->aid', new_tt_cores[-1],
                                       remainder)
  
        else:
          new_tt_cores[-1] = tf.einsum('oaib,obd->oaid', new_tt_cores[-1],
                                       remainder)
        remainder = None
      # TODO: infer the output ranks and shape.
      if do_collapse_batch_dim:
        return TensorTrain(new_tt_cores)
      else:
        return TensorTrainBatch(new_tt_cores)
  
    def __getitem__(self, slice_spec):
      """Basic indexing, returns a `TensorTrainBatch` with the specified region.
  
      Examples:
        >>> a = t3f.random_tensor_batch((2, 3, 4), batch_size=5)
        >>> a[1:3, :, :, :]
        is a 3D TensorTrainBatch 2 x 3 x 4 with batch_size = 2.
        >>> a[1:3]
        the same as above, a 3D TensorTrainBatch 2 x 3 x 4 with batch_size = 2.
        >>> a[1, :, :, :]
        is a 3D TensorTrain 2 x 3 x 4.
        >>> a[1]
        the same as above, a 3D TensorTrain 2 x 3 x 4.
        >>> a[1:3, :, 1, :]
        is a 2D TensorTrainBatch 2 x 4 with batch_size = 2.
        >>> a[1, :, 1, :]
        is a 2D TensorTrain 2 x 4.
  
      Returns:
        `TensorTrainBatch` or `TensorTrain` depending on whether the first
        (batch) dim was specified as a range or as a number.
      """
      try:
        slice_only_batch_dim = len(slice_spec) == 1
      except TypeError:
        # The argument is not iterable, so it's a single slice, or a number, or a
        # tf.Tensor with a number.
        slice_only_batch_dim = True
  
      if slice_only_batch_dim:
        # Indexing only for the batch_size axis, e.g. a[1:3].
        return self._batch_dim_getitem(slice_spec)
      elif len(slice_spec) == self.ndims() + 1:
        return self._full_getitem(slice_spec)
      else:
        raise ValueError('TensorTrainBatch.__getitem__: wrong number of '
                         'dimensions, expected 1 or %d, got %d' %
                         (self.ndims() + 1, len(slice_spec)))
  
  
  def _are_batch_tt_cores_valid(tt_cores, shape, tt_ranks, batch_size):
    """Check if dimensions of the TT-cores are consistent and the dtypes coincide.
  
    Args:
      tt_cores: a tuple of `Tensor` objects
      shape: An np.array, a tf.TensorShape (for tensors), a tuple of
        tf.TensorShapes (for TT-matrices or tensors), or None
      tt_ranks: An np.array or a tf.TensorShape of length len(tt_cores)+1.
      batch_size: a number or None
  
    Returns:
      boolean, True if the dimensions and dtypes are consistent.
    """
    shape = shapes.clean_raw_shape(shape)
    num_dims = len(tt_cores)
  
    for core_idx in range(1, num_dims):
      if tt_cores[core_idx].dtype != tt_cores[0].dtype:
        return False
    try:
      for core_idx in range(num_dims):
        curr_core_shape = tt_cores[core_idx].shape.as_list()
        if len(curr_core_shape) != len(tt_cores[0].get_shape()):
          # Shapes are inconsistent.
          return False
        if batch_size is not None and curr_core_shape[0] is not None:
          if curr_core_shape[0] != batch_size:
            # The TT-cores are not aligned with the given batch_size.
            return False
        if shape is not None:
          for i in range(len(shape)):
            dim_a, dim_b = curr_core_shape[i + 2], shape[i][core_idx]
            if dim_a is not None and dim_b is not None and dim_a != dim_b:
              # The TT-cores are not aligned with the given shape.
              return False
        if core_idx >= 1:
          prev_core_shape = tt_cores[core_idx - 1].get_shape()
          if curr_core_shape[1] != prev_core_shape[-1]:
            # TT-ranks are inconsistent.
            return False
        if tt_ranks is not None:
          if curr_core_shape[1] != tt_ranks[core_idx]:
            # The TT-ranks are not aligned with the TT-cores shape.
            return False
          if curr_core_shape[-1] != tt_ranks[core_idx + 1]:
            # The TT-ranks are not aligned with the TT-cores shape.
            return False
      if tt_cores[0].get_shape()[1] != 1 or tt_cores[-1].get_shape()[-1] != 1:
        # The first or the last rank is not 1.
        return False
    except ValueError:
      # The shape of the TT-cores is undetermined, can not validate it.
      pass
    return True
  
  
  def _infer_batch_raw_shape(tt_cores):
    """Tries to infer the (static) raw shape from the TT-cores."""
    num_dims = len(tt_cores)
    num_tensor_shapes = len(tt_cores[0].get_shape().as_list()) - 3
    raw_shape = [[] for _ in range(num_tensor_shapes)]
    for dim in range(num_dims):
      curr_core_shape = tt_cores[dim].get_shape()
      for i in range(num_tensor_shapes):
        raw_shape[i].append(curr_core_shape[i + 2])
    for i in range(num_tensor_shapes):
      raw_shape[i] = tf.TensorShape(raw_shape[i])
    return tuple(raw_shape)
  
  
  def _infer_batch_tt_ranks(tt_cores):
    """Tries to infer the (static) raw shape from the TT-cores."""
    tt_ranks = []
    for i in range(len(tt_cores)):
      tt_ranks.append(tt_cores[i].get_shape()[1])
    tt_ranks.append(tt_cores[-1].get_shape()[-1])
    return tf.TensorShape(tt_ranks)
  