tenstorrent / pytorch2.0_ttnn

⭐️ TTNN Compiler for PyTorch 2.0 ⭐️ It enables running PyTorch2.0 models on Tenstorrent hardware
https://tenstorrent.github.io/tt-metal/latest/ttnn/
25 stars 6 forks source link

1D shape is lost in conversion to `TILE_LAYOUT` tensor #280

Open jerrysky3 opened 1 month ago

jerrysky3 commented 1 month ago

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 https://github.com/tenstorrent/tt-metal/issues/12671 where we get 2D output on elementwise op with 1D input

jerrysky3 commented 1 month ago

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 commented 1 month ago

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