319d76de1d740a0500b9bc321ddc40bf3c032f71,tensorflow_transform/tf_utils.py,,reduce_batch_count_mean_and_var,#Any#Any#,646

Before Change

  if isinstance(x, tf.SparseTensor):
    // This means reduce_instance_dims=False.
    // TODO(b/112656428): Support SparseTensors with rank other than 2.
    if x.get_shape().ndims != 2:
      raise NotImplementedError(
          "Mean and var only support SparseTensors with rank 2")

    col_count, col_indices = x.dense_shape[1], x.indices[:, 1]
    x_sum = tf.math.unsorted_segment_sum(x.values, col_indices, col_count)
    x_mean = tf.where(tf.math.greater(x_count, 0),
                      x_sum / x_count,
                      tf.zeros_like(x_count, dtype=x.dtype))

    mean_values = tf.gather(x_mean, col_indices)
    x_minus_mean = x.values - mean_values

    col_sum_of_squares = tf.math.unsorted_segment_sum(
        tf.square(x_minus_mean), col_indices, col_count)

After Change

    x_mean = tf.where(tf.math.greater(x_count, 0),
                      x_sum / x_count,
                      tf.zeros_like(x_count, dtype=x.dtype))
    x_minus_mean = tf.sparse.add(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(

Instances