Linux-cpp-lisp / opt_einsum_fx

Einsum optimization using opt_einsum and PyTorch FX graph rewriting

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[idea] fuse einsums

mariogeiger opened this issue · comments

Make a transformer that detects einsums that can be fused.

Examples

These two einsum can be fuse together

a = einsum(...)
b = einsum(..., a)

but not these ones

a = einsum(...)
b = einsum(..., a)
c = einsum(..., a)

because a is used twice and we cannot ensure that repeating the operation is better.

This is now available as fuse_einsums (see tests in test_fuse.py)