Thank you @psrivas2 for your response! Sorry for coming back late on this but I was going through the source code to understand the pass better. I'm planning to add support for AXIS_SEPARATOR and have proposed below mentioned changes after going through the code. Can you please review them and share your feedback?
1. To update inputs with index_map, `TransformLayout` is invoked [here](https://github.com/apache/tvm/blob/unity/src/relax/transform/alter_op_impl.cc#L211). I plan to add AXIS_SEPARATOR in `LayoutTransformAttrs` to be saved for respective buffers. We can get the AXIS_SEPARATOR from user in the same way we are getting `index_map`. 2. Similarly for output buffers, `TransformLayout` gets invoked in [TransformLayoutInverse](https://github.com/apache/tvm/blob/unity/src/relax/transform/alter_op_impl.cc#L177) function. 3. The `replacement_primfunc` passed to `AlterOpImpl` pass by user will already have support for AXIS_SEPARATOR because it's a user configurable function. Please let me know if any point above needs any further clarification and I will be happy to answer. As always, thank you for your help! :slight_smile: --- [Visit Topic](https://discuss.tvm.apache.org/t/relax-alteroplayout-with-axis-separator/15113/3) to respond. You are receiving this because you enabled mailing list mode. To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/c1d95d57dd3447104afdb3470784c36ed05e3656534ab41ad9c433feb1ac6735).
