// This is to workaround the unnecessary cast i64 -> i32.
    SmallVector<int32_t, 4> new_shape_array;
    for (auto size : output_type.getShape()) {
      new_shape_array.push_back(ConvertToTfliteSize(size));
    }

    auto new_shape = rewriter.create<TFL::ConstOp>(
        loc, DenseIntElementsAttr::get(
                 tensorflow::GetTypeFromTFTensorShape(