def _reduce_batch_count_mean_and_var_ragged(
    x: tf.RaggedTensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
  """Computes elementwise count, mean and var for the given ragged tensor."""
  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=0)
  dense_x_count = tf.cast(
      reduce_batch_count(x, reduce_instance_dims=False), x.dtype)
  x_count = tf.RaggedTensor.from_tensor(
      dense_x_count, lengths=x_sum.nested_row_lengths())
  x_mean = tf.math.divide_no_nan(x_sum, x_count).to_tensor()
  dense_x = x.to_tensor()
  dense_x_is_finite = _is_finite(dense_x)
  x_minus_mean = tf.where(dense_x_is_finite, dense_x - x_mean,
                          tf.zeros_like(dense_x))
  x_minus_mean = tf.RaggedTensor.from_tensor(
      x_minus_mean, lengths=x.nested_row_lengths())
  sum_of_squares = tf.reduce_sum(input_tensor=tf.square(x_minus_mean), axis=0)
  x_variance = tf.math.divide_no_nan(sum_of_squares, x_count)
  return (dense_x_count, x_mean, x_variance.to_tensor())
Computes elementwise count, mean and var for the given ragged tensor.
tensorflow_transform/tf_utils.py