Fixed the bug to access tensor stride (#150)

Summary:
Pull Request resolved: #150

There is a bug to access tensor stride. The current implementation uses the index of the input tensor to access node.tensor_strides. However  node.tensor_strides is for all input variables including non tensor variables. So the fix is to get the list of tensor stride, then uses the index of the input tensor to access the list of tensor stride.

Differential Revision: D60642585

et_replay/execution_trace.py

@@ -292,9 +292,30 @@ def get_tensors(self, param_list: Iterable) -> List[tuple]:
                 tensors.extend(self.get_tensors(zip(elem_type, input, shape)))
         return tensors
 
+    def get_tensor_strides(
+        self, input_list: Iterable, stride_list: Iterable
+    ) -> List[tuple]:
+        strides = []
+        for (type, input, shape), stride in zip(input_list, stride_list):
+            if type.startswith("Tensor"):
+                strides.append(stride)
+            # GenericList could have tensor elements
+            elif type.startswith("GenericList"):
+                elem_type = type[12:-1].split(",")
+                strides.extend(
+                    self.get_tensor_strides(zip(elem_type, input, shape), stride)
+                )
+        return strides
+
     def get_input_tensors(self) -> List[tuple]:
         return self.get_tensors(self.get_inputs())
 
+    def get_input_tensor_strides(self) -> Optional[List[tuple]]:
+        if self.input_strides is None:
+            return None
+        else:
+            return self.get_tensor_strides(self.get_inputs(), self.input_strides)
+
     def get_output_tensors(self) -> List[tuple]:
         return self.get_tensors(self.get_outputs())
 

et_replay/tools/et_replay.py

@@ -517,6 +517,7 @@ def allocate_tensors(self):
                         self.args.pooling_factor,
                         self.args.alpha,
                     )
+            tensor_strides = node.get_input_tensor_strides()
             for idx, (data_type, t_id, shape) in enumerate(get_input_tensors(node)):
                 device = self.device
                 if self.tensor_with_device:
@@ -549,7 +550,7 @@ def allocate_tensors(self):
 
                             strides = None
                             if node.input_strides is not None:
-                                strides = node.input_strides[idx]
+                                strides = tensor_strides[idx]
                             tensor = self.get_tensor_from_storage(
                                 t_id[1],  # storage_id
                                 t_id[2],  # offset