triton.language.join

triton.language.join(a, b)

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

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

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

如果想合并两个以上的元素，可以多次调用此函数。这反映了Triton中张量尺寸必须是2的幂次方的约束条件。

join 是 split 的逆操作。

Parameters:

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

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