Fix DecomposeMultiDimSqueeze callback failure

[meenakshiramanathan1]

Fixes #642

[nvukobratTT]

@meenakshiramanathan1 can you provide a bit more context to this PR. Here are few questions I'm interested in:

1. The one @dgolubovicTT already asked, are we testing this path? If no, let's try to reproduce this with simple op test.
2. What are the results with and without your fix?
   • Before the change, you end up with failure in DecomposeMultiDimSqueeze
   • What about when you add this fix? Howe this multi-dim squeeze is further lowered in FFE, and latter on into MLIR? Can you share some reportify graphs for few compile stages (both TVM and FFE), as well a TTIR?
   • If you're unable to create simple unit test to cover this one, let's see screenshots from the original model.

[meenakshiramanathan1]

@meenakshiramanathan1 can you provide a bit more context to this PR. Here are few questions I'm interested in:

1. The one @dgolubovicTT already asked, are we testing this path? If no, let's try to reproduce this with simple op test.

2. What are the results with and without your fix?

   • Before the change, you end up with failure in DecomposeMultiDimSqueeze
   • What about when you add this fix? Howe this multi-dim squeeze is further lowered in FFE, and latter on into MLIR? Can you share some reportify graphs for few compile stages (both TVM and FFE), as well a TTIR?
   • If you're unable to create simple unit test to cover this one, let's see screenshots from the original model.

1. Tried testing this path with dynamic inputs but in the all cases of trying to create dynamic inputs , it is ending up with formation of argwhere which is not supported in forge yet
   For example,

def test_dynamic_multi_squeeze(test_device):
    class Multi_Squeeze(torch.nn.Module):
        def __init__(self):
            super().__init__()

        def forward(self, input):
            mask = input > 5
            indices = torch.nonzero(mask)
            squeezed_tensor = indices.squeeze(0)
            return squeezed_tensor

Here non_zero is mapped to argwhere which was one of the cases which was tried, so if we want to test this out further , argwhere op tvm decomposition should be brought in.

Before fix , graps are dumped till relay passes.

While after fix, it is progressing till forge_passes

2. With my fix , model is ending up with argwhere unsupported and without my fix, it would end up with this error.

3. Screenshots from original model which is causing dynamic shapes:

[nvukobratTT]

@meenakshiramanathan1

Tried testing this path with dynamic inputs but in the all cases of trying to create dynamic inputs , it is ending up with formation of argwhere which is not supported in forge yet

Thanks for the details! We're pushing the effort on outlining what and how our indexing should work.

Here non_zero is mapped to argwhere which was one of the cases which was tried, so if we want to test this out further , argwhere op tvm decomposition should be brought in.

Before fix , graps are dumped till relay passes.

Can you send actually screenshot of the graph? I'm interested in graph structure changes, rather than compile depth.

[meenakshiramanathan1]

@meenakshiramanathan1

Tried testing this path with dynamic inputs but in the all cases of trying to create dynamic inputs , it is ending up with formation of argwhere which is not supported in forge yet

Thanks for the details! We're pushing the effort on outlining what and how our indexing should work.

Here non_zero is mapped to argwhere which was one of the cases which was tried, so if we want to test this out further , argwhere op tvm decomposition should be brought in.
Before fix , graps are dumped till relay passes.

Can you send actually screenshot of the graph? I'm interested in graph structure changes, rather than compile depth.

I will send this information.

[meenakshiramanathan1]

@meenakshiramanathan1

Tried testing this path with dynamic inputs but in the all cases of trying to create dynamic inputs , it is ending up with formation of argwhere which is not supported in forge yet

Thanks for the details! We're pushing the effort on outlining what and how our indexing should work.

Here non_zero is mapped to argwhere which was one of the cases which was tried, so if we want to test this out further , argwhere op tvm decomposition should be brought in.
Before fix , graps are dumped till relay passes.

Can you send actually screenshot of the graph? I'm interested in graph structure changes, rather than compile depth.