BuildBodyWithCall(rewriter, loc, func, func_ty, &reduce_op.getBody());

    rewriter.replaceOp(op, reduce_op.getResults());

    return success();
  }
};

// Convert tf.XlaVariadicSort to mhlo.Sort
class ConvertXlaVariadicSortOp
    : public OpRewritePattern<TF::XlaVariadicSortOp> {
 public:
  using OpRewritePattern::OpRewritePattern;