auto strides = ArrayAttr::get(context, llvm::to_vector<4>(attrs));
  conv2d->setAttr("strides", strides);
}

// Transforms input shape for the first convolution.
void HandleConv2DInput(TF::Conv2DOp conv2d, int64_t block_size) {
  auto input = conv2d.getInput();
  auto input_shape = mlir::cast<RankedTensorType>(input.getType()).getShape();
  SmallVector<int64_t, 4> transform_shape = {