tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
Apache License 2.0
400 stars 50 forks source link

Need permute support for dimension > 4 #12154

Open punithsekar opened 2 weeks ago

punithsekar commented 2 weeks ago

Describe the bug Need ttnn support for 6D permute, used in swin_s model

To Reproduce Steps to reproduce the behavior:

  1. Run the below code snippet
@pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
def test_6d_permute(device):
    input_a=torch.randn(8,19,19,7,7,96,dtype=torch.bfloat16)
    tt_input_a=ttnn.from_torch(input_a,layout=ttnn.TILE_LAYOUT,device=device)
    tt_input_a=ttnn.permute(tt_input_a,(0,1,3,2,4,5))

Screenshots

 RuntimeError: TT_FATAL @ ../ttnn/cpp/ttnn/operations/data_movement/permute/permute.cpp:162: input_rank <= 4
**Please complete the following environment information:**
punithsekar commented 2 weeks ago

cc: @saichandax