import tensorflow as tf

          else:
            remainder = tf.matmul(remainder, sliced_core)
        else:
          if remainder is not None:
            # Add reminder from the previous collapsed cores to the current
            # core.
            sliced_core = tf.einsum('ab,bid->aid', remainder, sliced_core)
            remainder = None
          new_tt_cores.append(sliced_core)

    if remainder is not None:
      # The reminder obtained from collapsing the last cores.
      new_tt_cores[-1] = tf.einsum('aib,bd->aid', new_tt_cores[-1], remainder)
      remainder = None
    # TODO: infer the output ranks and shape.
    return TensorTrain(new_tt_cores)


def _are_tt_cores_valid(tt_cores, shape, tt_ranks):
  """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
