support dimension generalization for torch.Tensor.expand() (#383)

* support checking model redundancy

* revert change of vision_model_test

* reformat python code.

* reformat bert_model_test.py and utils.py

* minor fix

* fix failed check by comparing directories after os.path.realpath()

* fix bugs in check_validate.sh

* set dynamic=False in single_device_runner.py

* reset graph hash

* add robustness code for generating input tensor constraints

* Introduce input_tensor_constraints.py using shape propagation logic.

* support dimension generalization for torch.Tensor.view and torch.Tensor.reshape

* 1) support dimension generalization for torch.Tensor.expand(); 2) fix bugs in generalization for torch.Tensor.view and torch.Tensor.reshape

Author: lixinqi

graph_net/config/small_torch_samples_list.txt

@@ -1,3 +1,3 @@
-samples/timm/resnetaa50d.d_in12k
+# samples/timm/resnetaa50d.d_in12k
 # samples/transformers-auto-model/opus-mt-en-gmw
-# samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation
+samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation

graph_net/torch/static_to_dynamic.py

@@ -74,6 +74,7 @@ def call_method_predicator(method_name):
         return [
             (call_method_predicator("view"), dynamic_view_rewrite_pass),
             (call_method_predicator("reshape"), dynamic_reshape_rewrite_pass),
+            (call_method_predicator("expand"), dynamic_expand_rewrite_pass),
         ]
 
     def forward(self, *args, **kwargs):
@@ -87,6 +88,13 @@ def forward(self, *args, **kwargs):
         # return traced_module(*args, **kwargs)
 
 
+def has_call_method(traced_module: fx.GraphModule, method_name) -> bool:
+    for node in traced_module.graph.nodes:
+        if node.op == "call_method" and node.target == method_name:
+            return True
+    return False
+
+
 def dynamic_view_rewrite_pass(traced_module: fx.GraphModule) -> fx.GraphModule:
     """
     Fx Pass: Replaces hardcoded constants in 'view' ops that match an input tensor dimension
@@ -141,12 +149,15 @@ def dynamic_view_rewrite_pass(traced_module: fx.GraphModule) -> fx.GraphModule:
                     # and this dimension is one we wish to generalize (e.g., batch size, axis 0)
                     if target_dim == input_dim_size:
                         # Prioritize matching the batch size (axis 0)
-                        if axis_idx == 0:
-                            matched_axis = 0
+                        if i == 0 and axis_idx == 0:
+                            matched_axis = axis_idx
                             break
-                        # If not axis 0, but still a match, consider generalization
-                        elif matched_axis == -1:
+                        elif i > 0 and axis_idx > 0 and input_dim_size > 1:
                             matched_axis = axis_idx
+                            break
+                        else:
+                            # Do nothing.
+                            pass
 
                 if matched_axis != -1:
                     # Found a matching dynamic axis (matched_axis), replace it with a size() call
@@ -242,12 +253,15 @@ def dynamic_reshape_rewrite_pass(traced_module: fx.GraphModule) -> fx.GraphModul
                     # and this dimension is one we wish to generalize (e.g., batch size, axis 0)
                     if target_dim == input_dim_size:
                         # Prioritize matching the batch size (axis 0)
-                        if axis_idx == 0:
-                            matched_axis = 0
+                        if i == 0 and axis_idx == 0:
+                            matched_axis = axis_idx
                             break
-                        # If not axis 0, but still a match, consider generalization
-                        elif matched_axis == -1:
+                        elif i > 0 and axis_idx > 0 and input_dim_size > 1:
                             matched_axis = axis_idx
+                            break
+                        else:
+                            # Do nothing.
+                            pass
 
                 if matched_axis != -1:
                     # Found a matching dynamic axis (matched_axis), replace it with a size() call
@@ -289,8 +303,100 @@ def dynamic_reshape_rewrite_pass(traced_module: fx.GraphModule) -> fx.GraphModul
     return traced_module
 
 
-def has_call_method(traced_module: fx.GraphModule, method_name) -> bool:
+def dynamic_expand_rewrite_pass(traced_module: fx.GraphModule) -> fx.GraphModule:
+    """
+    Fx Pass: Replaces hardcoded constants in 'expand' ops that match an input tensor dimension
+    with a dynamic 'size()' call. The primary goal is to dynamicize the batch size (axis 0).
+    """
+    # Create a new graph to hold the rewritten nodes
+    new_graph = fx.Graph()
+
+    # Create a map to link nodes from the old graph to nodes in the new graph
+    val_map = {}
+
     for node in traced_module.graph.nodes:
-        if node.op == "call_method" and node.target == method_name:
-            return True
-    return False
+        if node.op == "call_method" and node.target == "expand":
+            # Get the input tensor node
+            input_tensor_node = node.args[0]
+            # Get the target shape arguments for expand (e.g., 1, 4, 6, 64)
+            expand_args = node.args[1:]
+
+            # --- Dependency on ShapeProp Results ---
+            # input_shape is the static shape (e.g., batch_size, C, H, W)
+            input_meta = input_tensor_node.meta.get("tensor_meta")
+            if input_meta is None:
+                raise RuntimeError(
+                    f"Node {input_tensor_node.name} lacks tensor_meta. Did ShapeProp run?"
+                )
+
+            input_shape = input_meta.shape
+
+            # Find the new list of expand arguments
+            new_expand_args = []
+
+            # Iterate over the target dimensions of expand (dim0, dim1, ...)
+            for i, target_dim in enumerate(expand_args):
+                # 1. Handle dynamic dimensions (e.g., -1 or non-integer values)
+                if not isinstance(target_dim, int) or target_dim < 1:
+                    new_expand_args.append(
+                        val_map[target_dim] if target_dim in val_map else target_dim
+                    )
+                    continue
+
+                # 2. Handle hardcoded constants (e.g., 1, 6, 64)
+
+                # --- Core Logic: Find the matching dynamic axis ---
+
+                # Default: Keep the hardcoded constant if no matching dynamic axis is found
+                best_match = target_dim
+                matched_axis = -1
+
+                axis_idx = i
+                input_dim_size = input_shape[i]
+                if target_dim == input_dim_size:
+                    if axis_idx == 0:
+                        matched_axis = axis_idx
+                    elif axis_idx > 0 and input_dim_size > 1:
+                        matched_axis = axis_idx
+                    else:
+                        # Do nothing.
+                        pass
+
+                if matched_axis != -1:
+                    # Found a matching dynamic axis (matched_axis), replace it with a size() call
+
+                    # 1. Create a call to size(axis) in the new graph
+                    # NOTE: input_tensor_node must first be mapped to a new graph node via val_map
+                    new_input_node = val_map[input_tensor_node]
+
+                    # Use the size() method to retrieve the dynamic dimension
+                    size_node = new_graph.call_method(
+                        "size", args=(new_input_node, matched_axis)
+                    )
+
+                    best_match = size_node
+
+                new_expand_args.append(best_match)
+
+            # --- Rebuild the expand node ---
+            # 1. Map the input tensor node to the new graph node
+            new_input_node = val_map[input_tensor_node]
+
+            # 2. Insert the new expand node into the new graph
+            # with new_graph.inserting_after(new_input_node):
+            new_node = new_graph.call_method(
+                "expand", args=(new_input_node, *new_expand_args)
+            )
+
+            # 3. Map the old node to the new node
+            val_map[node] = new_node
+
+        else:
+            # Copy other nodes to the new graph
+            new_node = new_graph.node_copy(node, lambda x: val_map[x])
+            val_map[node] = new_node
+
+    # Replace the old graph with the new graph and return
+    traced_module.graph = new_graph
+    traced_module.recompile()
+    return traced_module