def _reduce_batch_count_mean_and_var_dense(
    x: tf.Tensor,
    reduce_instance_dims: bool) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
  """Computes count, mean and var for the given dense tensor."""
  axis = None if reduce_instance_dims else 0
  x_count = tf.cast(reduce_batch_count(x, reduce_instance_dims), x.dtype)
  zeros_like_x = tf.zeros_like(x)
  x_is_finite = _is_finite(x)
  x_sum = tf.reduce_sum(tf.where(x_is_finite, x, zeros_like_x), axis=axis)
  x_mean = tf.math.divide_no_nan(x_sum, x_count)
  x_minus_mean = tf.where(x_is_finite, x - x_mean, zeros_like_x)
  sum_of_squares = tf.reduce_sum(
      input_tensor=tf.square(x_minus_mean), axis=axis)
  x_variance = tf.math.divide_no_nan(sum_of_squares, x_count)
  return (x_count, x_mean, x_variance)
Computes count, mean and var for the given dense tensor.
tensorflow_transform/tf_utils.py