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

Before Change

      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))

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)

Instances