// Only the first convolution in the graph whose batch size smaller than 8
// and its input feature size smaller than 8 can be transformed.
std::optional<BlockArgumentInfo> GetConv2DInputArgNum(TF::Conv2DOp conv2d) {
  if (conv2d.getDataFormat() != "NHWC" || conv2d.getStrides().size() != 4) {
    return std::nullopt;
  }
  // Current supported ops between convolution input and the block arguments are