triton.language.reshape

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

Returns a tensor with the same number of elements as input but with the provided shape.

Parameters:
  • input (Block) – The input tensor.

  • shape – The new shape.

shape can be passed as a tuple or as individual parameters:

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

This function can also be called as a member function on tensor, as x.reshape(...) instead of reshape(x, ...).