triton.language.split

triton.language.split(a) tuple[张量, 张量]

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

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

如果您想要分割成超过两块,可以多次调用此函数(可能还需要配合reshape操作)。这反映了Triton中张量尺寸必须是2的幂次方的约束条件。

split 是 join 的逆操作。

Parameters:

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

该函数也可以作为成员函数在tensor上调用, 使用x.split()而不是 split(x)