// In the following, we'll check indices is obtained by a iota.
    auto sort_dim_attr = DenseIntElementsAttr::get(
        RankedTensorType::get({1}, rewriter.getI64Type()), {sort_dim});
    if (!MatchIota(sort_dim_attr, indices))
      return rewriter.notifyMatchFailure(
          op, "the second operand is supposed to be obtained from IOTA");
    if (!MatchTopKComparator<mhlo::ReturnOp>(op.getComparator()))