auto new_shape = rewriter.create<TF::ShapeOp>(loc, shape_type, input);
    SmallVector<int64_t, 8> output_shape(/*Size=*/1, op.getNumElements());
    for (const auto &dim : dense_elem_attr.getValues<APInt>())
      output_shape.push_back(dim.getSExtValue());
    RankedTensorType result_type = tensorflow::GetTypeFromTFTensorShape(
        output_shape, getElementTypeOrSelf(input));
    rewriter.replaceOpWithNewOp<TF::ReshapeOp>(op, result_type, input,