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.
  // We will reshape the weight into [output, 1, 1, depth]