Function ArgMaxTwoPass

tensorflow/compiler/xla/client/lib/arithmetic.cc:224–226

}

XlaOp ArgMaxTwoPass(XlaOp input, PrimitiveType output_type, int axis) {
  return ArgMinMaxTwoPass(input, output_type, axis, /*is_min=*/false);
}

XlaOp ArgMinTwoPass(XlaOp input, PrimitiveType output_type, int axis) {
  return ArgMinMaxTwoPass(input, output_type, axis, /*is_min=*/true);