triton.language.split

triton.language.split(a) → tuple[tensor, tensor]

沿张量的最后一个维度将其分成两半，该维度的大小必须为 2。

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

如果您想分成多于两部分，可以使用多次调用此函数（可能再加上调用 reshape）。这反映了 Triton 中张量大小必须是 2 的幂的约束。

split 是 join 的逆操作。

参数:

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

此函数也可以作为 tensor 的成员函数调用，写成 x.split() 而不是 split(x)。