ViT Support

[cnut1648]

Hello, am I right that currently TE does not support ViT? Specifically, I am looking for ViT that uses 2D rope, e.g. Pixtral's ViT.

I also checked Nemo codebase, but it seems that they are also building their own impl for ViT instead of relying on TE.

Nemo defines ParallelVisionTransformer. This class init a series of ParallelVisionTransformerLayer layers, each a subclass of ParallelVisionTransformerLayer_, which in turn a subclass of ParallelTransformerLayer_. This class defines all operations in Nemo's impl.

Thanks.