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,