def reduce_batch_count(x: common_types.TensorType,
                       reduce_instance_dims: bool) -> tf.Tensor:
  """Counts elements in the given tensor.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    reduce_instance_dims: A bool, if True - collapses the batch and instance
      dimensions to arrive at a single scalar output. Otherwise, only collapses
      the batch dimension and outputs a `Tensor` of the same shape as the input.

  Returns:
    The element count of `x`. The result is either a scalar if
    reduce_instance_dims is True, otherwise a `Tensor` having shape of `x`
    without the first (batch) dimension. NaNs and infinite input values are
    ignored.
  """
  if isinstance(x, tf.SparseTensor):
    if reduce_instance_dims:
      x = x.values
    else:
      ones_like = tf.SparseTensor(
          indices=x.indices,
          values=tf.cast(_is_finite(x.values), tf.int64),
          dense_shape=x.dense_shape)
      # TODO(b/178189903): Remove this once we no longer lose static shape
      # information.
      ones_like._dense_shape_default = x._dense_shape_default  # pylint: disable=protected-access
      return _sparse_reduce_batch_keep_shape(tf.sparse.reduce_sum, ones_like)
  elif isinstance(x, tf.RaggedTensor):
    if reduce_instance_dims:
      x = x.flat_values
    else:
      finite_mask = tf.cast(_is_finite(x), tf.int64)
      return tf.math.reduce_sum(finite_mask, axis=0).to_tensor()

  # Exlude NaNs and infinite elements from size calculation. They can only occur
  # in tensors with floating data types.
  if x.dtype.is_floating:
    finite_mask = tf.cast(tf.math.is_finite(x), tf.int64)
    return tf.reduce_sum(finite_mask, axis=None if reduce_instance_dims else 0)

  if reduce_instance_dims:
    return tf.size(input=x)

  # Fill a tensor shaped like x except batch_size=1 with batch_size.
  x_shape = tf.shape(input=x)
  return tf.fill(x_shape[1:], x_shape[0])
Counts elements in the given tensor.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    reduce_instance_dims: A bool, if True - collapses the batch and instance
      dimensions to arrive at a single scalar output. Otherwise, only collapses
      the batch dimension and outputs a `Tensor` of the same shape as the input.

  Returns:
    The element count of `x`. The result is either a scalar if
    reduce_instance_dims is True, otherwise a `Tensor` having shape of `x`
    without the first (batch) dimension. NaNs and infinite input values are
    ignored.
  
tensorflow_transform/tf_utils.py