opt_einsum_fx:使用opt_einsum和PyTorch FX图形重写进行Einsum优化 源码

opt_einsum_fx 正在开发中 使用opt_einsum和PyTorch FX图形重写进行Einsum优化。 该库当前支持: 将多个枚举融合为一个 使用库优化einsum 用标量常量融合乘法和除法,包括通过与einsum之类的运算进行融合,这些操作与标量乘法相通。 安装 $ pip install 您可以使用以下命令运行测试 $ pytest tests/ 用法 opt_einsum_fx基于 ,该框架是用于在PyTorch Python代码和以编程方式可操纵的计算图之间进行转换的框架。 要使用此程序包,必须有可能将您的函数或模型作为torch.fx.Graph :讨论了FX符号跟踪的局限性。 最小的例子 import torch import torch . fx import opt_einsum_fx def einmatvecmul ( a , b , vec )