1D shape is lost in conversion to TILE_LAYOUT tensor

[jerrysky3]

When converting an 1D tensor into TILE_LAYOUT, the result tensor will have the shape (1[32], x) and the original 1D shape information is lost during the process. When we want to convert the tiled tensor back to pytorch tensor, currently I don't know how to tell if the original shape is 1D.

For example, the code below is doing a round-trip conversion between pytorch and ttnn tensor and I expect to get the output (32,) but the result is (1, 32)

import torch
import ttnn


def main(device):
    input = torch.rand((32,), dtype=torch.bfloat16)
    input_tensor = ttnn.from_torch(input, layout=ttnn.TILE_LAYOUT, device=device)
    print(ttnn.to_torch(input_tensor).shape)


if __name__ == "__main__":
    device = ttnn.open_device(device_id=0)
    try:
        main(device)
    finally:
        ttnn.close_device(device)

I think this might be the root cause of the issues like tenstorrent/tt-metal#12671 where we get 2D output on elementwise op with 1D input

[jerrysky3]

For the completeness of the conversion, maybe we need to store the actual original 1D shape in TTNN tensor object so to_torch can convert it back to the orignal shape?

[davorchap]

For the completeness of the conversion, maybe we need to store the actual original 1D shape in TTNN tensor object so to_torch can convert it back to the orignal shape?

Yes! @TT-BrianLiu and @ayerofieiev-tt are making the new amazing shape/layout in Tensor