def word_count(tokens: Union[tf.SparseTensor, tf.RaggedTensor],
               name: Optional[str] = None) -> tf.Tensor:
  # pyformat: disable
  """Find the token count of each document/row.

  `tokens` is either a `RaggedTensor` or `SparseTensor`, representing tokenized
  strings. This function simply returns size of each row, so the dtype is not
  constrained to string.

  Example:
  >>> sparse = tf.SparseTensor(indices=[[0, 0], [0, 1], [2, 2]],
  ...                          values=['a', 'b', 'c'], dense_shape=(4, 4))
  >>> tft.word_count(sparse)
  <tf.Tensor: shape=(4,), dtype=int64, numpy=array([2, 0, 1, 0])>

  Args:
    tokens: either
      (1) a `SparseTensor`, or
      (2) a `RaggedTensor` with ragged rank of 1, non-ragged rank of 1
      of dtype `tf.string` containing tokens to be counted
    name: (Optional) A name for this operation.

  Returns:
    A one-dimensional `Tensor` the token counts of each row.

  Raises:
    ValueError: if tokens is neither sparse nor ragged
  """
  # pyformat: enable
  with tf.compat.v1.name_scope(name, 'word_count'):
    if isinstance(tokens, tf.RaggedTensor):
      return tokens.row_lengths()
    elif isinstance(tokens, tf.SparseTensor):
      result = tf.sparse.reduce_sum(
          tf.SparseTensor(indices=tokens.indices,
                          values=tf.ones_like(tokens.values, dtype=tf.int64),
                          dense_shape=tokens.dense_shape),
          axis=list(range(1, tokens.get_shape().ndims)))
      result.set_shape([tokens.shape[0]])
      return result
    else:
      raise ValueError('Invalid token tensor')
Find the token count of each document/row.

  `tokens` is either a `RaggedTensor` or `SparseTensor`, representing tokenized
  strings. This function simply returns size of each row, so the dtype is not
  constrained to string.

  Example:
  >>> sparse = tf.SparseTensor(indices=[[0, 0], [0, 1], [2, 2]],
  ...                          values=['a', 'b', 'c'], dense_shape=(4, 4))
  >>> tft.word_count(sparse)
  <tf.Tensor: shape=(4,), dtype=int64, numpy=array([2, 0, 1, 0])>

  Args:
    tokens: either
      (1) a `SparseTensor`, or
      (2) a `RaggedTensor` with ragged rank of 1, non-ragged rank of 1
      of dtype `tf.string` containing tokens to be counted
    name: (Optional) A name for this operation.

  Returns:
    A one-dimensional `Tensor` the token counts of each row.

  Raises:
    ValueError: if tokens is neither sparse nor ragged
  
tensorflow_transform/mappers.py