EinsumPass

Description

Constructs a computational graph based on the expression of the Einsum operator.

In the static shape scenario, there are two cases:

Static shape and single input:

Before: After:

Static shape and dual inputs (Depending on the specific expression, the two inputs of the MatMul operator may be reversed):

After:

In the dynamic shape scenario, the graph to which the Einsum operator belongs before fusion has dual inputs, as shown in the following figure:

Currently, 14 types of expressions are supported in the dynamic shape scenario. The figures below show the graphs after fusion.

The Einsum expression is "abc,cde->abde".
The Einsum expression is "abcd,aecd->aceb". (The two inputs of the BatchMatMul operator are reversed.)
The Einsum expression is "abcd,adbe->acbe".
The Einsum expression is "abcd,cde->abe".
The Einsum expression is "abc,cd->abd".
The Einsum expression is "abc,dc->abd".
The Einsum expression is "abc,abd->dc". (The two inputs of the MatMulV2 operator are reversed.)
The Einsum expression is "abc,dec->abde".
The Einsum expression is "abc,abde->dec".
Tensor B with a static shape:

Tensor B with a dynamic shape:
The Einsum expression is "abcd,aecd->acbe".
The Einsum expression is "abcd,acbe->aecd".
The Einsum expression is "abcd,ecd->abe".
The Einsum expression is "abcd,abe->ecd".
Tensor A with a static shape:

Tensor A with a dynamic shape:
The Einsum expression is "abcd,acbe->adbe".

Based on whether the input shape is dynamic, the Reshape dym/stc node contains the operators shown in the following figure.

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.

The Einsum expression is "abcd, abde->abce".
The Einsum expression is "abcd, abce->abde".

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.
The Einsum expression is "abcd, aebd->aebc".

Based on whether the input shape is dynamic, the Reshape node contains the operators shown in the following figure.
The Einsum expression is "abcd, abce->acde".

Based on whether the input shape is dynamic, the Reshape node contains the operators shown in the following figure.
The Einsum expression is "abc, abd->acd".

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.
The Einsum expression is "ab, cb->ac".

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.
The Einsum expression is "abc, acd->abd".
The Einsum expression is "abc, adc->abd".

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.
The Einsum expression is "abcd, ced->abce".

Based on whether the input shape is dynamic, the Reshape node contains the operators shown in the following figure.

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.
The Einsum expression is "abcd, ebcd->bcae".

Based on whether the input shape is dynamic, the Reshape node contains the operators shown in the following figure.

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.
The Einsum expression is "abcd, dabe->cabe".
Based on whether the input shape is dynamic, the Reshape node contains the operators shown in the following figure.
The Einsum expression is "a, b->ab".
The Einsum expression is "abcd, aecd->eb".
Based on whether the input shape is dynamic, the Reshape node contains the operators shown in the following figure.

Based on whether the input shape is dynamic, the Transpose dym/stc node contains the operators shown in the following figure.
The Einsum expression is "abcd, eb->aecd":

Based on whether the input shape is dynamic, the Reshape node contains the operators shown in the following figure.

In other dynamic shape scenarios, there are two cases:
1. Single-input scenario:
  Before: After:
2. Dual-input scenario:
  1. Scenario 1
    Before: After:
    
    Unsqueeze, ReduceSum, and Squeeze are optional operators, and each of them may appear multiple times.
  2. Scenario 2
    Before: After:
    
    Unsqueeze, ReduceSum, Unsqueeze/FlattenV2, Reshape, Transpose, and Squeeze are optional operators, and ReduceSum and Unsqueeze/FlattenV2 may appear multiple times.

Restrictions

This fusion pattern is enabled by default and cannot be disabled.
The Einsum operator can have only one or two inputs, and the number of elements of the input tensor cannot overflow the int64 representation range.
In the static shape scenario, the Einsum expression supports only a single input with a single output or dual inputs with a single output.
In the static shape scenario, the input or output expression cannot have duplicate dimension labels.
In the static shape scenario, if the fused operator is BatchMatMul, the data type of the input tensor can only be float16 or float32.
In the dynamic shape scenario, only 14 Einsum expressions are supported (the specific dimension label is only an example): "abc,cde->abde", "abcd,aecd->aceb", "abcd,adbe->acbe", "abcd,cde->abe", "abc,cd->abd", "abc,dc->abd", "abc,abd->dc", "abc,dec->abde", "abc,abde->dec", "abcd,aecd->acbe", "abcd,acbe->aecd", "abcd,ecd->abe", "abcd,abe->ecd", and "abcd,acbe->adbe". If the Einsum operator has only one input, the input cannot have a dynamic shape.
In the dynamic shape scenario, the number of input dimensions must be the same as that in the corresponding Einsum input expression.
In the dynamic shape scenario, the Einsum expression supports a single input with a single output, dual inputs with a single output, or empty output. The input or output expression cannot have duplicate dimension labels. A dimension label of the output must exist in the input.
In the dynamic shape scenario, only input tensors with a maximum of four dimensions are supported.

Availability

Atlas A2 training product / Atlas A2 inference product

Atlas 350 Accelerator Card

Parent topic: Graph Fusion Patterns