Support for einops

[escorciav]

Is there any plan for supporting einops?
While preparing a model from machine learning research scientis, I noticed that AIMET struggles with it while ONNX does not.

Is it due to fx?

Thanks!

[quic-klhsieh]

Hi @escorciav , we currently do not have plans to support automatically handling models using einops package during model preparation. You are correct that the limitation stems from torch.fx incompatibililty with einops.

In the past, we have gotten around this by using torch.fx.wrap() on einops functions in the model definition file. Otherwise, you will need to hand modify the model to use conventional pytorch operations.