Location loc = op.getLoc();
    ElementsAttr dimension;
    matchPattern(op.getDimension(), m_Constant(&dimension));
    // Create the mhlo.sort op.
    auto sort_op = rewriter.create<mhlo::SortOp>(
        loc, op.getInputs(), dimension.getValues<IntegerAttr>()[0].getInt(),
        op.getIsStable());
    mlir::SymbolRefAttr func = op.getComparator();