<commit_before>import numpy as np
import tensorflow as tf


# TODO: test all these functions.
def tt_ranks(tt):
  """Returns the TT-ranks of a TensorTrain.

  This operation returns a 1-D integer tensor representing the TT-ranks of
  the input.

  Args:
    tt: `TensorTrain` object.

  Returns:
    A `Tensor`
  """
  num_dims = tt.ndims()
  ranks = []
  for i in range(num_dims):
    ranks.append(tf.shape(tt.tt_cores[i])[0])
  ranks.append(tf.shape(tt.tt_cores[-1])[-1])
  return tf.stack(ranks, axis=0)


def shape(tt):
  """Returns the shape of a TensorTrain.

  This operation returns a 1-D integer tensor representing the shape of
  the input. For TT-matrices the shape would have two values, see raw_shape for
  the tensor shape.

  Args:
    tt: `TensorTrain` object.

  Returns:
    A `Tensor`
  """
  tt_raw_shape = raw_shape(tt)
  if tt.is_tt_matrix():
    return tf.reduce_prod(raw_shape, axis=1)
  else:
    return tt_raw_shape[0]
<commit_msg>Support batch TT in shapes.tt_ranks<commit_after>import numpy as np
import tensorflow as tf


# TODO: test all these functions.
def tt_ranks(tt):
  """Returns the TT-ranks of a TensorTrain.

  This operation returns a 1-D integer tensor representing the TT-ranks of
  the input.

  Args:
    tt: `TensorTrain` or `TensorTrainBatch` object.

  Returns:
    A `Tensor`
  """
  num_dims = tt.ndims()
  ranks = []
  # TODO: ugly.
  from tensor_train_batch import TensorTrainBatch
  left_rank_dim = 1 if isinstance(tt, TensorTrainBatch) else 0
  for i in range(num_dims):
    ranks.append(tf.shape(tt.tt_cores[i])[left_rank_dim])
  ranks.append(tf.shape(tt.tt_cores[-1])[-1])
  return tf.stack(ranks, axis=0)


def shape(tt):
  """Returns the shape of a TensorTrain.

  This operation returns a 1-D integer tensor representing the shape of
  the input. For TT-matrices the shape would have two values, see raw_shape for
  the tensor shape.

  Args:
    tt: `TensorTrain` object.

  Returns:
    A `Tensor`
  """
  tt_raw_shape = raw_shape(tt)
  if tt.is_tt_matrix():
    return tf.reduce_prod(raw_shape, axis=1)
  else:
    return tt_raw_shape[0]

{'fraction_non_alphanumeric': 0.10481181515007146, 'fraction_numerical': 0.008099094807050978, 'mean_word_length': 3.5258620689655173, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1299813', 'n_tokens_mistral': 763, 'n_tokens_neox': 728, 'n_words': 229}