Sweet! That is ideal. I'm gonna build a bunch of neural CRFs.
One more small thing. This line has me a bit worried:
EINSUM_SYMBOLS_BASE = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
Just want to be sure that this will all support arbitrary size CUDA variables. Seems like opt-einsum does that: