triton.language.split

triton.language.split(a, _semantic=None, _generator=None) → tuple[tensor, tensor]

将张量沿其最后一个维度（大小必须为 2）分割为两部分。

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

如果您想将其分割成两部分以上，可以使用对该函数的多次调用（可能需要加上调用 reshape）。这反映了 Triton 中张量必须具有 2 的幂次大小的限制。

split 是 join 的逆操作。

参数:

a (Tensor) – 要分割的张量。

此函数也可以作为成员函数在 tensor 上调用，形式为 x.split() 而非 split(x)。