bug fix

Author: KRRT7

code_to_optimize/code_directories/simple_tracer_e2e/workload.py

@@ -21,6 +21,45 @@ def test_threadpool() -> None:
     for r in result:
         print(r)
 
+class AlexNet:
+    def __init__(self, num_classes=1000):
+        self.num_classes = num_classes
+        self.features_size = 256 * 6 * 6
+
+    def forward(self, x):
+        features = self._extract_features(x)
+        
+        output = self._classify(features)
+        return output
+
+    def _extract_features(self, x):
+        result = []
+        for i in range(len(x)):
+            pass
+
+        return result
+
+    def _classify(self, features):
+        total = sum(features)
+        return [total % self.num_classes for _ in features]
+
+class SimpleModel:
+    @staticmethod
+    def predict(data):
+        return [x * 2 for x in data]
+    
+    @classmethod
+    def create_default(cls):
+        return cls()
+
+def test_models():
+    model = AlexNet(num_classes=10)
+    input_data = [1, 2, 3, 4, 5]
+    result = model.forward(input_data)
+
+    model2 = SimpleModel.create_default()
+    prediction = model2.predict(input_data)
 
 if __name__ == "__main__":
     test_threadpool()
+    test_models()

codeflash/discovery/discover_unit_tests.py

@@ -195,17 +195,25 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                 if node.module == "importlib" and alias.name == "import_module":
                     self.has_dynamic_imports = True
 
-                # Check if imported name is a target qualified name
-                if alias.name in self.function_names_to_find:
-                    self.found_any_target_function = True
-                    self.found_qualified_name = alias.name
-                    return
-                # Check if module.name forms a target qualified name
-                qualified_name = f"{node.module}.{alias.name}"
-                if qualified_name in self.function_names_to_find:
-                    self.found_any_target_function = True
-                    self.found_qualified_name = qualified_name
-                    return
+            qualified_name = f"{node.module}.{alias.name}"
+            potential_matches = {alias.name, qualified_name}
+
+            if any(name in self.function_names_to_find for name in potential_matches):
+                self.found_any_target_function = True
+                self.found_qualified_name = next(
+                    name for name in potential_matches if name in self.function_names_to_find
+                )
+                return
+
+            qualified_prefix = qualified_name + "."
+            if any(target_func.startswith(qualified_prefix) for target_func in self.function_names_to_find):
+                self.found_any_target_function = True
+                self.found_qualified_name = next(
+                    target_func
+                    for target_func in self.function_names_to_find
+                    if target_func.startswith(qualified_prefix)
+                )
+                return
 
     def visit_Attribute(self, node: ast.Attribute) -> None:
         """Handle attribute access like module.function_name."""

tests/scripts/end_to_end_test_utilities.py

@@ -206,8 +206,8 @@ def run_trace_test(cwd: pathlib.Path, config: TestConfig, expected_improvement_p
         return False
 
     functions_traced = re.search(r"Traced (\d+) function calls successfully and replay test created at - (.*)$", stdout)
-    if not functions_traced or int(functions_traced.group(1)) != 4:
-        logging.error("Expected 4 traced functions")
+    if not functions_traced or int(functions_traced.group(1)) != 13:
+        logging.error("Expected 13 traced functions")
         return False
 
     replay_test_path = pathlib.Path(functions_traced.group(2))