tf.math.count_nonzero