Remove fallback to 2-pass algorithm for argmin/argmax on TPU. (#3831)
Remove fallback to 2-pass algorithm for argmin/argmax on TPU. (#3831)

(The compiler problem that prompted the workaround seems to be fixed.)
google/jax,google/jax,google/jax,tensorflow/probability,tensorflow/probability,google/jax