def reduce_batch_count_mean_and_var(
    x: common_types.TensorType,
    reduce_instance_dims: bool) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
  """Computes element count, mean and var for 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:
    A 3-tuple containing the tensor's (count, mean, var). NaNs and infinite
    input values are ignored.
  """
  if isinstance(x, tf.SparseTensor):
    if reduce_instance_dims:
      return _reduce_batch_count_mean_and_var_dense(
          x.values, reduce_instance_dims=True)
    else:
      return _reduce_batch_count_mean_and_var_sparse(x)
  elif isinstance(x, tf.RaggedTensor):
    if reduce_instance_dims:
      return _reduce_batch_count_mean_and_var_dense(
          x.flat_values, reduce_instance_dims=True)
    else:
      return _reduce_batch_count_mean_and_var_ragged(x)
  else:
    return _reduce_batch_count_mean_and_var_dense(x, reduce_instance_dims)
Computes element count, mean and var for 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:
    A 3-tuple containing the tensor's (count, mean, var). NaNs and infinite
    input values are ignored.
  
tensorflow_transform/tf_utils.py