search

bdusell / semiring-einsum

Generic PyTorch implementation of einsum that supports different semirings
https://bdusell.github.io/semiring-einsum/
MIT License
44 stars 8 forks source link

Avoid the call to permute in LookupInfo.lookup #33

Open ccshan opened 1 year ago

ccshan commented 1 year ago

Call unsqueeze and permute ahead of time in LookupInfo.view, then just index into the result in view_lookup instead of lookup

This speeds up https://github.com/diprism/fggs/blob/main/test/test_sum_product.py by ~10%

bdusell commented 1 year ago

Thanks! I'll take a look and try to get this merged soon.