Output shape = ops::Add(s.WithOpName("shape"), shape_a, shape_b);

  Output reshape_input = ops::Placeholder(s.WithOpName("reshape_input"),
                                          DT_FLOAT, ops::Placeholder::Attrs{});
  Output reshape = ops::Reshape(s.WithOpName("reshape"), reshape_input, shape);

  AddToCluster({shape.node(), reshape.node()}, "cluster_0");

                            ops::Const(root.WithOpName("maxval"), 20));
  Output reshape_input =
      ops::Placeholder(root.WithOpName("reshape_input"), DT_FLOAT,
                       ops::Placeholder::Shape(TensorShape({500, 500})));
  Output reshape =
      ops::Reshape(root.WithOpName("reshape"), reshape_input, shape);

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));

    if (!reshape_op) return failure();

    // Check if the last dimension does not change after reshape.
    auto reshape_input = reshape_op.getInput();
    auto reshape_input_ty = mlir::dyn_cast<ShapedType>(reshape_input.getType());
    if (!reshape_input_ty.hasStaticShape() || input_ty.getRank() == 0 ||
        reshape_input_ty.getRank() == 0 ||

  const int64_t flat_size = input_type.getNumElements();
  const int64_t width = flat_size / depth;
  SmallVector<int64_t, 4> input_new_shape({1, 1, width, depth});
  auto reshaped_input =
      InsertReshapeOp(fc_op.getLoc(), input, input_type.getElementType(),
                      input_new_shape, &rewriter);

  // Insert a reshape after the weight.

                     Type element_type) {
  Value shape_cst = BuildIntArrayConstOp(builder, rewriter, shape, idx_type);
  Value reshaped_input = builder.create<TF::ReshapeOp>(
      RankedTensorType::get(shape, element_type), input, shape_cst);
  return reshaped_input;
}

// Create a tensor which is equal to input[begin: begin + size].
Value BuildSliceOp(ImplicitLocOpBuilder& builder,

    Value zero_matrix = rewriter.create<BroadcastOp>(
        op.getLoc(), reshaped_input.getType(), zero,
        GetI64ElementsAttr({new_size, new_size}, &rewriter));
    Value masked =
        rewriter.create<SelectOp>(op.getLoc(), reshaped_input.getType(),
                                  compare, reshaped_input, zero_matrix);
    auto reduce = rewriter.create<ReduceOp>(op.getLoc(), masked, zero,