@yzhliu Gave it some more thoughts over last few days. I think there might be slightly better way to deal with layouts.
* Instead of directly using the `transpose` operators, maybe we can use some new annotation ops like `annotate.change_layout` (or maybe use `layout_transform`). This will have a src and dst layout. For each op that goes through legalization, we can have one `change_layout` in the start and one at the end. * Now, instead of `propagating` the layout, I would hoist/sink the `change_layout` annotation ops if the ops are layout agnostic. This is similar to Loop Invariant Code Motion (or here, if you allow me, Layout Invariant Code Motion ;-) ). We can hoist these annotated ops as much as possible. Therefore, you would not insert any new ops, you will just try to move them to right place. * If done properly, the end state will have back to back `change_layout` ops. Now, a peephole can remove them. And then a realize pass can convert the remaining annotated ops to a sequence of existing Relay ops. This should be able to work for both TF and MxNet, i.e., support both NCHW and NHWC layout. -- You are receiving this because you are subscribed to this thread. Reply to this email directly or view it on GitHub: https://github.com/dmlc/tvm/issues/3670#issuecomment-522068048