def squeeze_batch_dim(tt):
  """Converts batch size 1 TensorTrainBatch into TensorTrain.

  Args:
    tt: TensorTrain or TensorTrainBatch.

  Returns:
    TensorTrain if the input is a TensorTrainBatch with batch_size == 1 (known
      at compilation stage) or a TensorTrain.
    TensorTrainBatch otherwise.
    """
  try:
    if tt.batch_size == 1:
      return tt[0]
    else:
      return tt
  except AttributeError:
    # tt object does not have attribute batch_size, probably already
    # a TensorTrain.
    return tt