For space-to-depth ops, instead of using the memory-copy-reshape based ops provided by frameworks such as TensorFlow [1]} and Pytorch [2]}, we customize an \(n \times n\)  convolution with stride-n to perform the space-to-depth operation, reshaping an \(H\times W \times C \)  tensor to an \(H/n \times W/n \times C*n^2 \)  tensor. This approach has two advantages: 1) convolutions are much preferred by TPUs and GPUs because of their high operational intensity and execution efficiency; 2) in addition to reshaping the input tensor to improve operational intensity and efficiency, the \(n \times n\)  convolutions can also be trained to contribute to the model's capacity. For space-to-batch ops, we have to use the memory-intensive copy-reshape ops provided by common frameworks [1]}, [2]}.
