def reduce_batch_count_mean_and_var_per_key(
    x: common_types.TensorType, key: common_types.TensorType,
    reduce_instance_dims: bool
) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
  """Computes per-key element count, mean and var for the given tensor.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    key: A `Tensor` or `CompositeTensor` (cannot be None).
        Must meet one of the following conditions:
        1. Both x and key are dense,
        2. Both x and key are composite and `key` must exactly match `x` in
        everything except values,
        3. The axis=1 index of each element of sparse x matches its index of
        dense key.
    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. Not supported for `CompositeTensor`s.

  Returns:
    A 4-tuple containing the `Tensor`s (key_vocab, count, mean, var). NaNs and
    infinite input values are ignored.
  """

  if isinstance(x, (tf.SparseTensor, tf.RaggedTensor)):
    if not reduce_instance_dims:
      raise NotImplementedError(
          'Mean and var per key only support reduced dims for CompositeTensors')

  x, key = _validate_and_get_dense_value_key_inputs(x, key)

  unique = tf.unique(key, out_idx=tf.int64)
  x_is_finite = _is_finite(x)

  finite_x = tf.where(x_is_finite, x, tf.zeros_like(x))
  if reduce_instance_dims:
    x_count = tf.cast(x_is_finite, x.dtype)
    if x.get_shape().ndims != 1:
      x_count = tf.reduce_sum(x_count, axis=1)
    x_count = tf.math.unsorted_segment_sum(x_count, unique.idx,
                                           tf.size(unique.y))
    sums = (
        tf.reduce_sum(finite_x, axis=1)
        if x.get_shape().ndims != 1 else finite_x)
    sums = tf.math.unsorted_segment_sum(sums, unique.idx, tf.size(unique.y))
  else:
    sums = tf.math.unsorted_segment_sum(finite_x, unique.idx, tf.size(unique.y))
    x_count = tf.math.unsorted_segment_sum(
        tf.cast(x_is_finite, tf.float32), unique.idx, tf.size(unique.y))

  means = tf.math.divide_no_nan(tf.cast(sums, x.dtype), x_count)
  sum_sqs = tf.math.unsorted_segment_sum(
      tf.square(finite_x), unique.idx, tf.size(input=unique.y))
  if sum_sqs.get_shape().ndims != 1 and reduce_instance_dims:
    sum_sqs = tf.reduce_sum(sum_sqs, axis=1)

  variances = tf.math.divide_no_nan(sum_sqs, x_count) - tf.square(means)

  return unique.y, tf.cast(x_count, tf.int64), means, variances
Computes per-key element count, mean and var for the given tensor.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    key: A `Tensor` or `CompositeTensor` (cannot be None).
        Must meet one of the following conditions:
        1. Both x and key are dense,
        2. Both x and key are composite and `key` must exactly match `x` in
        everything except values,
        3. The axis=1 index of each element of sparse x matches its index of
        dense key.
    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. Not supported for `CompositeTensor`s.

  Returns:
    A 4-tuple containing the `Tensor`s (key_vocab, count, mean, var). NaNs and
    infinite input values are ignored.
  
tensorflow_transform/tf_utils.py