Make sure meta val of aten.index is correct

Author: swimdi

torch_ttnn/passes/lowering/to_tt_pass.py

@@ -17,6 +17,7 @@
 
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 import torch.fx.traceback as fx_traceback
+from torch._subclasses.fake_tensor import FakeTensorMode
 from . import target_wrappers
 from .to_tt_guard import can_lowering_to_ttnn
 from operator import getitem
@@ -447,7 +448,7 @@ def __init__(self, node):
 
     def call_function(self, target, args=(), kwargs={}):
         new_node = self.g.call_function(target, args, kwargs)
-        new_node.meta = self.node.meta
+        new_node.meta = self.node.meta.copy()
         if hasattr(self.node.target, "_schema"):
             new_node.meta["original_input_variations"] = metrics.collect_input_variation_from_node(self.node)
         if target == ttnn.layer_norm:
@@ -1133,16 +1134,21 @@ def batch_norm_inference(input, weight, bias, mean, var, momentum, eps):
 
 
 def DigestAtenOps(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
-    g = gm.graph
     nodes = list(gm.graph.nodes)
     for node in nodes:
+        g = GraphWrapper(node)
 
         def rewrite_node(node):
             args = node.args
             kwargs = node.kwargs
 
             if node.target == torch.ops.aten.index.Tensor:
 
+                def edit_meta_val(node, shape, dtype):
+                    fake_mode = FakeTensorMode()
+                    fake_tensor = fake_mode.from_tensor(torch.zeros(shape, dtype=dtype))
+                    node.meta["val"] = fake_tensor
+
                 def broadcast_indices(indices):
                     import numpy as np
 
@@ -1156,27 +1162,47 @@ def broadcast_indices(indices):
                             broadcasted_indices.append(
                                 g.call_function(torch.ops.aten.expand.default, (indices[i], broadcasted_shape))
                             )
+                            edit_meta_val(broadcasted_indices[-1], broadcasted_shape, indices[i].meta["val"].dtype)
                     return broadcasted_shape, broadcasted_indices
 
                 # for example, input.shape = (3, 4, 5), indices = [tensor([[0, 1, 1]]), tensor([[2, 1, 2]])]
                 # then output is [[input[0][2], input[1][1], input[1][2]]]
                 input_tensor, indices = args
+                if get_shape(input_tensor) is None:
+                    return None
+                if None in [get_shape(indices[i]) for i in range(len(indices))]:
+                    return None
                 index_shape, indices = broadcast_indices(indices)
                 input_shape = get_shape(input_tensor)
                 num_index = len(indices)
                 index_size = index_shape.numel()
                 remained_shape = input_shape[num_index:]
                 reshape_shape = index_shape + remained_shape
-                indices_flatten = [g.call_function(torch.ops.aten.flatten, args=(idx,)) for idx in indices]
+                input_dtype = input_tensor.meta["val"].dtype
+                flatten_shape = torch.Size([index_size])
+                indices_flatten = [
+                    g.call_function(torch.ops.aten.reshape.default, args=(idx, flatten_shape)) for idx in indices
+                ]
+                for i in range(len(indices_flatten)):
+                    edit_meta_val(indices_flatten[i], flatten_shape, indices[i].meta["val"].dtype)
                 output = []
                 for i in range(index_size):
                     indexing = [g.call_function(getitem, args=(indices_flatten[n], i)) for n in range(num_index)]
+                    for n in range(num_index):
+                        edit_meta_val(indexing[n], torch.Size([]), indices_flatten[n].meta["val"].dtype)
                     output.append(g.call_function(getitem, args=(input_tensor, indexing)))
+                    edit_meta_val(output[-1], remained_shape, input_dtype)
                 # aten.cat cannot concat zero dim tensor
                 if len(remained_shape) == 0:
-                    output = [g.call_function(torch.ops.aten.reshape, args=(o, [1])) for o in output]
-                output_cat = g.call_function(torch.ops.aten.cat, args=(output,))
-                output_reshape = g.call_function(torch.ops.aten.reshape, args=(output_cat, reshape_shape))
+                    remained_shape = torch.Size([1])
+                    output = [g.call_function(torch.ops.aten.reshape.default, args=(o, remained_shape)) for o in output]
+                    for o in output:
+                        edit_meta_val(o, remained_shape, input_dtype)
+                output_cat = g.call_function(torch.ops.aten.cat.default, args=(output,))
+                output_cat_shape = torch.Size([len(output)] + list(remained_shape))
+                edit_meta_val(output_cat, output_cat_shape, input_dtype)
+                output_reshape = g.call_function(torch.ops.aten.reshape.default, args=(output_cat, reshape_shape))
+                edit_meta_val(output_reshape, reshape_shape, input_dtype)
                 return output_reshape
 
         with g.inserting_before(node):