Missing and inconsistent einsum operations by tahashieenavaz · Pull Request #363 · lucidrains/vit-pytorch

tahashieenavaz · 2026-04-28T23:07:08Z

Following:

Fixes missing einsum operations
Makes them consistently be imported from torch which is faster and not einops .
Adds smoke tests to simple_flash_attn_vit_3d and simple_flash_attn_vit.

tahashieenavaz added 8 commits April 29, 2026 00:51

adds missing einsum imports

127a7a1

adds simple test to simple_flash_attn_vit

9047e7a

adds smoke test to simple flash attn vit 3d

960e66d

imports einsum directly for consistency

6f348e0

replaces einsum from torch with einsum from einops

7889a82

replace einsum from torch with einsum from einops

5651393

merges torch imports

c6e8352

replaces einsum

045a640

Navigation Menu