first version with tracin

Author: misrasaurabh1

codeflash/discovery/discover_unit_tests.py

@@ -322,6 +322,7 @@ def filter_test_files_by_imports(
     """
     if not target_functions:
         return file_to_test_map
+    return file_to_test_map # added as a temporary fix, can be removed
 
     logger.debug(f"Target functions for import filtering: {target_functions}")
 

codeflash/model.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+
+class AlexNet(nn.Module):
+
+    def __init__(self, num_classes: int=1000, dropout: float=0.5) -> None:
+        super().__init__()
+        self.features = nn.Sequential(nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(64, 192, kernel_size=5, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(192, 384, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2))
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
+        self.classifier = nn.Sequential(nn.Dropout(p=dropout), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=False), nn.Dropout(p=dropout), nn.Linear(4096, 4096), nn.ReLU(inplace=False), nn.Linear(4096, num_classes))
+
+    def classifier_forward(self, x: torch.Tensor):
+        return self.classifier(x)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return self.classifier_forward(x)

codeflash/model_run.py

@@ -0,0 +1,37 @@
+from PIL import Image
+import urllib
+from torchvision import transforms
+import torch
+
+from codeflash.tracer import Tracer
+from codeflash.model import AlexNet
+
+def load_model():
+    model = AlexNet()
+    model.load_state_dict(torch.load('alexnet-owt-7be5be79.pth'))
+    model.eval()
+    return model
+
+def download_image():
+    url, filename = ('https://github.com/pytorch/hub/raw/master/images/dog.jpg', 'dog.jpg')
+    try:
+        urllib.URLopener().retrieve(url, filename)
+    except:
+        urllib.request.urlretrieve(url, filename)
+
+download_image()
+model = load_model()
+input_image = Image.open('dog.jpg')
+preprocess = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+input_tensor = preprocess(input_image)
+input_batch = input_tensor.unsqueeze(0)
+if torch.cuda.is_available():
+    input_batch = input_batch.to('cuda')
+    model.to('cuda')
+    print('YESS')
+with Tracer():
+    with torch.no_grad():
+        output = model(input_batch)
+print(output[0].shape)
+probabilities = torch.nn.functional.softmax(output[0], dim=0)
+print(probabilities.shape)

codeflash/optimization/function_optimizer.py

@@ -908,7 +908,7 @@ def generate_tests_and_optimizations(
             function_to_concolic_tests, concolic_test_str = future_concolic_tests.result()
             logger.info(f"Generated {len(tests)} tests for {self.function_to_optimize.function_name}")
             console.rule()
-            generated_tests = GeneratedTestsList(generated_tests=tests)
+            generated_tests = GeneratedTestsList(generated_tests=[])
 
         return Success(
             (