code: 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")
  
imports: import os
  from typing import Any, Callable, Iterable, Optional, Tuple, Union
  import tensorflow
  from tensorflow_transform import analyzers
  from tensorflow_transform import common
  from tensorflow_transform import common_types
  from tensorflow_transform import gaussianization
  from tensorflow_transform import schema_inference
  from tensorflow_transform import tf_utils
  from tfx_bsl.types import tfx_namedtuple
label_desc: Use image node_id 28 for calling a global function with example usage: word_count(tokens, name) and returns: tokens, result