Merge branch 'main' into fix/verify-existing-optimizations-error-handling

Author: KRRT7

code_to_optimize/code_directories/simple_tracer_e2e/codeflash.trace

@@ -1,74 +1,39 @@
 from concurrent.futures import ThreadPoolExecutor
-from time import sleep
 
 
 def funcA(number):
-    number = number if number < 1000 else 1000
+    number = number if number < 100 else 100
     k = 0
-    for i in range(number * 100):
+    for i in range(number * 10):
         k += i
-    # Simplify the for loop by using sum with a range object
     j = sum(range(number))
-
-    # Use a generator expression directly in join for more efficiency
     return " ".join(str(i) for i in range(number))
 
 
 def test_threadpool() -> None:
-    pool = ThreadPoolExecutor(max_workers=3)
-    args = list(range(10, 31, 10))
+    pool = ThreadPoolExecutor(max_workers=2)
+    args = [5, 10, 15]
     result = pool.map(funcA, args)
 
     for r in result:
         print(r)
 
 class AlexNet:
-    def __init__(self, num_classes=1000):
+    def __init__(self, num_classes=10):
         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):
-        result = []
-        sleep(0.1) # can be optimized away
-        for i in range(500):
-            for x in data:
-                computation = 0
-                computation += x * i ** 2
-                result.append(computation)
-        return result
-    
-    @classmethod
-    def create_default(cls):
-        return cls()
+        result = 0
+        for val in x:
+            result += val * val
+        return result % self.num_classes
 
 
 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()

code_to_optimize/code_directories/simple_tracer_e2e/workload.py

@@ -0,0 +1,100 @@
+"""
+Utility functions used in CompEcon
+
+Based routines found in the CompEcon toolbox by Miranda and Fackler.
+
+References
+----------
+Miranda, Mario J, and Paul L Fackler. Applied Computational Economics
+and Finance, MIT Press, 2002.
+
+"""
+from functools import reduce
+import numpy as np
+import torch
+
+def _gridmake2(x1, x2):
+    """
+    Expands two vectors (or matrices) into a matrix where rows span the
+    cartesian product of combinations of the input arrays. Each column of the
+    input arrays will correspond to one column of the output matrix.
+
+    Parameters
+    ----------
+    x1 : np.ndarray
+         First vector to be expanded.
+
+    x2 : np.ndarray
+         Second vector to be expanded.
+
+    Returns
+    -------
+    out : np.ndarray
+          The cartesian product of combinations of the input arrays.
+
+    Notes
+    -----
+    Based of original function ``gridmake2`` in CompEcon toolbox by
+    Miranda and Fackler.
+
+    References
+    ----------
+    Miranda, Mario J, and Paul L Fackler. Applied Computational Economics
+    and Finance, MIT Press, 2002.
+
+    """
+    if x1.ndim == 1 and x2.ndim == 1:
+        return np.column_stack([np.tile(x1, x2.shape[0]),
+                               np.repeat(x2, x1.shape[0])])
+    elif x1.ndim > 1 and x2.ndim == 1:
+        first = np.tile(x1, (x2.shape[0], 1))
+        second = np.repeat(x2, x1.shape[0])
+        return np.column_stack([first, second])
+    else:
+        raise NotImplementedError("Come back here")
+
+
+def _gridmake2_torch(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+    """
+    PyTorch version of _gridmake2.
+
+    Expands two tensors into a matrix where rows span the cartesian product
+    of combinations of the input tensors. Each column of the input tensors
+    will correspond to one column of the output matrix.
+
+    Parameters
+    ----------
+    x1 : torch.Tensor
+         First tensor to be expanded.
+
+    x2 : torch.Tensor
+         Second tensor to be expanded.
+
+    Returns
+    -------
+    out : torch.Tensor
+          The cartesian product of combinations of the input tensors.
+
+    Notes
+    -----
+    Based on original function ``gridmake2`` in CompEcon toolbox by
+    Miranda and Fackler.
+
+    References
+    ----------
+    Miranda, Mario J, and Paul L Fackler. Applied Computational Economics
+    and Finance, MIT Press, 2002.
+
+    """
+    if x1.dim() == 1 and x2.dim() == 1:
+        # tile x1 by x2.shape[0] times, repeat_interleave x2 by x1.shape[0]
+        first = x1.tile(x2.shape[0])
+        second = x2.repeat_interleave(x1.shape[0])
+        return torch.column_stack([first, second])
+    elif x1.dim() > 1 and x2.dim() == 1:
+        # tile x1 along first dimension
+        first = x1.tile(x2.shape[0], 1)
+        second = x2.repeat_interleave(x1.shape[0])
+        return torch.column_stack([first, second])
+    else:
+        raise NotImplementedError("Come back here")