ttnn.reshape_on_device

ttnn.reshape_on_device() → None

Returns a tensor with the new shape of [W, Z, Y, X]. The X dimension of input and output tensor must have same size.

Equivalent pytorch code:

input_tensor = torch.arange(4.)
W = 1
Z = 1
Y = 2
X = 2
output_tensor = torch.reshape(input_tensor, (W, Z, Y, X))

:param * input_tensor: Input Tensor. :param * W: W dim of tensor. :param * Z: Z dim of tensor. :param * Y: Y dim of tensor. :param * X: X dim of tensor.

:keyword * memory_config: Memory Config of the output tensor

Example

>>> tensor = ttnn.from_torch(torch.tensor((1, 4), dtype=torch.bfloat16), device=device)
>>> output = ttnn.reshape(tensor, 1, 1, 2, 2)