def _reduce_batch_count_mean_and_var_sparse(
    x: tf.SparseTensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
  """Computes elementwise count, mean and var for the given sparse tensor."""
  x_count = tf.cast(reduce_batch_count(x, reduce_instance_dims=False), x.dtype)
  finite_x = tf.SparseTensor(
      indices=x.indices,
      values=tf.where(_is_finite(x.values), x.values, tf.zeros_like(x.values)),
      dense_shape=x.dense_shape)
  x_sum = _sparse_reduce_batch_keep_shape(tf.sparse.reduce_sum, finite_x)
  x_mean = tf.math.divide_no_nan(x_sum, x_count)
  x_minus_mean = tf.sparse.add(finite_x, -tf.broadcast_to(x_mean, tf.shape(x)))
  x_minus_mean_sparse = tf.SparseTensor(x.indices,
                                        tf.gather_nd(x_minus_mean, x.indices),
                                        x.dense_shape)
  sum_of_squares = tf.math.reduce_sum(
      tf.square(tf.sparse.to_dense(x_minus_mean_sparse)), axis=0)
  x_variance = tf.math.divide_no_nan(sum_of_squares, x_count)
  return (x_count, x_mean, x_variance)
Computes elementwise count, mean and var for the given sparse tensor.
tensorflow_transform/tf_utils.py