triton.language.join

triton.language.join(a, b, _semantic=None)

在新的次要维度中连接给定的张量。

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

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

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

join 是 split 的逆操作。

参数:

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

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