Lowering to ttnn.embedding produces a different shape if the input is 1-D

[kevinwuTT]

This example is from the perceiver_io model.

In this model, the aten graph has this code:

import torch
import ttnn

arg5_1 = torch.rand((2048, 768), dtype=torch.bfloat16)
arange = torch.ops.aten.arange.start(0, 2048, device = "cpu", pin_memory = False)
embedding_1 = torch.ops.aten.embedding.default(arg5_1, arange)
print(embedding_1.shape)

The arange tensor is 1D with shape of (2048). The size of the output should be:

torch.Size([2048, 768])

This is lowered to the equivalent code below:

with ttnn.manage_device(device_id=0) as device:
    ttnn_from_torch_2 = ttnn.from_torch(arange, layout = ttnn.TILE_LAYOUT, device = device, dtype = ttnn.uint32)
    ttnn_from_device_2 = ttnn.from_device(ttnn_from_torch_2, )
    ttnn_to_layout_3 = ttnn.to_layout(ttnn_from_device_2, ttnn.ROW_MAJOR_LAYOUT, )
    ttnn_to_device_2 = ttnn.to_device(ttnn_to_layout_3, device = device)
    ttnn_from_torch_3 = ttnn.from_torch(arg5_1, layout = ttnn.TILE_LAYOUT, device = device, dtype = ttnn.bfloat16)
    ttnn_from_device_3 = ttnn.from_device(ttnn_from_torch_3, )
    ttnn_to_layout_4 = ttnn.to_layout(ttnn_from_device_3, ttnn.ROW_MAJOR_LAYOUT, )
    ttnn_to_device_3 = ttnn.to_device(ttnn_to_layout_4, device = device)
    ttnn_embedding_1 = ttnn.embedding(ttnn_to_device_2, ttnn_to_device_3, layout = ttnn.ROW_MAJOR_LAYOUT)
    ttnn_embedding_1 = ttnn.to_torch(ttnn_embedding_1)
    print(ttnn_embedding_1.shape)

However, the shape becomes:

torch.Size([1, 2048, 768])

The issue comes from the first ttnn.from_torch which changes the shape of the arange tensor, (2048), to (1[32], 2048) due to tilization. Changing the tensor back to ROW_MAJOR will not remove the extra dimension. The most straightforward solution will be to set layout to ROW_MAJOR in the first ttnn.from_torch call and removing the extra device/layout changes. A workaround would be to add a call to reshape or squeeze if arg[0] of ttnn.embedding has only a rank of 1. We can remove this workaround once we have the better fix.