diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp index 29da32cd1791c..815806f06b472 100644 --- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -1356,7 +1356,7 @@ RankedTensorType GatherOp::inferResultType(RankedTensorType sourceType, SmallVector resultShape(indicesType.getShape().drop_back()); resultShape.reserve(resultShape.size() + sourceType.getRank()); for (int64_t idx : llvm::seq(0, sourceType.getRank())) { - if (std::binary_search(gatherDims.begin(), gatherDims.end(), idx)) { + if (llvm::binary_search(gatherDims, idx)) { if (!rankReduced) resultShape.push_back(1); continue;