[](Type ty) { return mlir::cast<ShapedType>(ty).getElementType(); })};

    // Create the mhlo.reduce op.
    auto reduce_op = rewriter.create<mhlo::ReduceOp>(
        loc, op.getInputs(), op.getInitValues(),
        GetI64ElementsAttr(op.getDimensionsToReduce()), elementTypes);

    // Insert a call to the reducer in the region of the mhlo op.
    BuildBodyWithCall(rewriter, loc, func, func_ty, &reduce_op.getBody());