triton.language.interleave

triton.language.interleave(a, b)

将两个张量的值沿其最后一个维度进行交错。这两个张量必须具有相同的形状。等同于 tl.join(a, b).reshape(a.shape[:-1] + [2 * a.shape[-1]])

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

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