def reduce_batch_count_per_key(
    key: common_types.TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
  """Computes per-key counts in the given tensor.

  Args:
    key: A `Tensor` or `CompositeTensor`.

  Returns:
    A 2-tuple containing the tensor's (key_vocab, count_per_key).
  """
  key = _to_string(key)

  if isinstance(key, tf.SparseTensor):
    key = key.values
  elif isinstance(key, tf.RaggedTensor):
    key = key.flat_values
  key.set_shape([None])
  unique = tf.unique_with_counts(key, out_idx=tf.int64)

  return unique.y, unique.count
Computes per-key counts in the given tensor.

  Args:
    key: A `Tensor` or `CompositeTensor`.

  Returns:
    A 2-tuple containing the tensor's (key_vocab, count_per_key).
  
tensorflow_transform/tf_utils.py