triton.language.reshape

triton.language.reshape(input, *shape, can_reorder=False, _semantic=None, _generator=None)

返回一个张量,其元素数量与输入相同,但具有提供的形状。

参数:
  • input (Block) – 输入张量。

  • shape – 新的形状。

shape 可以作为元组传递,也可以作为单独的参数传递。

# These are equivalent
reshape(x, (32, 32))
reshape(x, 32, 32)

此函数也可以作为 tensor 的成员函数调用,例如 x.reshape(...) 而不是 reshape(x, ...)