}
  }
}

// Gets block size that is equal to stride from spatial dimension
// from convolution.
// Space to depth transform won't be triggered if block size <= 1.
int32_t GetConv2DBlockSize(TF::Conv2DOp conv2d) {
  SmallVector<int32_t, 4> strides(4, 1);
  for (int i = 0; i < 3; ++i) {
    strides[i] = mlir::cast<mlir::IntegerAttr>(conv2d.getStrides()[i]).getInt();
  }