Einsum

根据 equation 参数指定的维度，计算输入元素沿维度的总和，该参数以爱因斯坦求和约定编写。

• equation 为输入中的每个维度指定 ASCII 小写字母，字母顺序与维度顺序相同，每个输入之间用逗号分隔。

• equation 表示为 term1,term2…->output-term，其中每个 term 对应一个操作数张量，term 内的字符对应操作数维度。

• 用相同下标标记的维度必须匹配。

• 在多个输入中重复一个标签意味着这些轴将被相乘。

• 从输出中省略一个标签意味着将累加沿这些轴的值。

• 输出下标必须至少出现一次在某个输入操作数中，并且最多出现一次在输出中。

• 在隐式模式下，即如果 equation 不包含 ->，则表达式中出现一次的索引将按字母顺序递增成为输出的一部分。

• 在显式模式下，可以通过添加箭头 (->) 后跟输出的下标标签来控制输出。例如，ij,jk->ik 等同于 ij,jk。

• 空字符串 (“”) 对于标量操作数有效。

• equation 可能在不同元素（下标、箭头和逗号）之间包含空格 (SPC- 0x20)。

属性

equation 一个字符串，表示以爱因斯坦求和约定编写的求和方程。

输入

inputs: 类型为 T 的张量。最多可以设置两个输入。

输出

output: 类型为 T 的张量。

数据类型

T: float16, float32, bfloat16

形状信息

示例

Einsum

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 3))
layer = network.add_einsum(inputs=[in1], equation="ij->ji")
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[-3.0, -2.0, -1.0], [0.0, 1.0, 2.0]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape

expected[layer.get_output(0).name] = np.array(
    [
        [-3.0, 0.0],
        [-2.0, 1.0],
        [-1.0, 2.0],
    ]
)

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 3))
in2 = network.add_input("input2", dtype=trt.float32, shape=(3, 1))
layer = network.add_einsum(inputs=[in1, in2], equation="ik,kj->ij")
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[-3.0, -2.0, -1.0], [0.0, 1.0, 2.0]])

inputs[in2.name] = np.array([[1.0], [2.0], [3.0]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape

expected[layer.get_output(0).name] = np.array([[-10.0], [8.0]])

C++ API

有关 C++ IEinsumLayer 算子的更多信息，请参阅 C++ IEinsumLayer 文档。

Python API

有关 Python IEinsumLayer 算子的更多信息，请参阅 Python IEinsumLayer 文档。