// 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();
    auto func_op = cast<mlir::func::FuncOp>(SymbolTable::lookupSymbolIn(
        op->getParentOfType<mlir::ModuleOp>(), func));