triton.language.join

triton.language.join(a, b)

将给定的张量在新的次要维度上连接起来。

例如，给定两个形状为 (4,8) 的张量，生成一个形状为 (4,8,2) 的新张量。给定两个标量，返回一个形状为 (2) 的张量。

两个输入会被广播到相同的形状。

如果你想连接两个以上的元素，可以使用多次调用此函数。这反映了 Triton 中的约束，即张量的大小必须是 2 的幂。

join 是 split 的逆操作。

参数:

• a (张量) – 第一个输入张量。

• b (张量) – 第二个输入张量。