Fix a race condition in tf.where when it counts a certain number of true elements

from a reference Tensor, but that tensor's value changes before it writes their
indices.
Change: 124730221