architecture folder docstring update

poet/architectures/bert.py

@@ -21,6 +21,16 @@
 
 # Look-up ignored for now
 def make_transformer(SEQ_LEN, HIDDEN_DIM, I, HEADS, layers):
+    """
+    Appends a transformer block to the list of layers.
+
+    Parameters:
+        SEQ_LEN (int): Length of the sequence.
+        HIDDEN_DIM (int): Dimension of the hidden layers.
+        I (int): Intermediate dimension used in calculations.
+        HEADS (int): Number of attention heads.
+        layers (list): List of layers to which the transformer block will be appended.
+    """
     input_layer = layers[-1]
     layers.append(QueryKeyValueMatrix(SEQ_LEN, HIDDEN_DIM, I, HEADS, layers[-1]))  # Calculate Query
     layers.append(QKTMatrix(SEQ_LEN=SEQ_LEN, HIDDEN_DIM=I, I=SEQ_LEN, ATTN_HEADS=HEADS, input=layers[-1]))  # QK^T
@@ -36,6 +46,19 @@ def make_transformer(SEQ_LEN, HIDDEN_DIM, I, HEADS, layers):
 
 
 def BERTBase(SEQ_LEN=512, HIDDEN_DIM=768, I=64, HEADS=12, NUM_TRANSFORMER_BLOCKS=12):
+    """
+    Constructs a BERT model with specified parameters.
+
+    Parameters:
+        SEQ_LEN (int, optional): Length of the input sequence. Default is 512.
+        HIDDEN_DIM (int, optional): Dimension of the hidden layers. Default is 768.
+        I (int, optional): Intermediate dimension used in calculations. Default is 64.
+        HEADS (int, optional): Number of attention heads. Default is 12.
+        NUM_TRANSFORMER_BLOCKS (int, optional): Number of transformer blocks. Default is 12.
+
+    Returns:
+        list: List of model layers.
+    """
     layers = [InputLayer((SEQ_LEN, HIDDEN_DIM))]
     for _transformer in range(0, NUM_TRANSFORMER_BLOCKS):
         make_transformer(SEQ_LEN, HIDDEN_DIM, I, HEADS, layers)

poet/architectures/linear.py

@@ -2,6 +2,12 @@
 
 
 def make_linear_network(batch_size=1, input_shape=[784]):
+    """
+    Constructs a linear network model with specified input shape and batch size.
+
+    Returns:
+        list: List of model layers.
+    """
     layers = [InputLayer((batch_size, *input_shape))]
 
     linear_layers = [[784, 10], [10, 120], [120, 100], [100, 200], [200, 10], [10, 10]]
@@ -18,6 +24,12 @@ def make_linear_network(batch_size=1, input_shape=[784]):
 
 
 def make_unit_linear_network(nfwd=12):
+    """
+    Constructs a unit linear network model with a specified number of forward layers.
+
+    Returns:
+        list: List of model layers.
+    """
     layers = [InputLayer((1,))]
     for i in range(nfwd):
         layers.append(ReLULayer(layers[-1]))

poet/architectures/resnet.py

@@ -17,7 +17,32 @@
 
 
 def resnet18(batch_size, num_classes=1000, input_shape=(3, 224, 224)):  # Imagenet
+    """
+    Constructs a ResNet-18 model.
+
+    Parameters:
+        batch_size (int): Size of the input batch.
+        num_classes (int, optional): Number of output classes. Default is 1000.
+        input_shape (tuple, optional): Shape of the input data. Default is (3, 224, 224).
+
+    Returns:
+        list: List of model layers.
+    """
+
     def make_basic_block(in_planes, planes, stride, padding, x):
+        """
+        Constructs a basic block for the ResNet-18 model.
+
+        Parameters:
+            in_planes (int): Number of input channels.
+            planes (int): Number of output channels.
+            stride (int): Stride of the convolutional layer.
+            padding (tuple): Padding for the convolutional layer.
+            x (Tensor): Input tensor.
+
+        Returns:
+            list: List of basic block layers.
+        """
         kernel_size = (3, 3)
         conv1 = Conv2dLayer(in_planes, planes, kernel_size, stride, padding, x)
         bn1 = BatchNorm2d(conv1)

poet/architectures/vgg.py

@@ -13,7 +13,30 @@
 
 
 def vgg16(batch_size, num_classes=1000, input_shape=(3, 224, 224)):
+    """
+    Constructs a VGG-16 model.
+
+    Parameters:
+        batch_size (int): Size of the input batch.
+        num_classes (int, optional): Number of output classes. Default is 1000.
+        input_shape (tuple, optional): Shape of the input data. Default is (3, 224, 224).
+
+    Returns:
+        list: List of model layers.
+    """
     def make_conv_stack(in_channels, out_filters, kernel_size, x):
+        """
+        Constructs a convolutional stack for the VGG-16 model.
+
+        Parameters:
+            in_channels (int): Number of input channels.
+            out_filters (int): Number of output filters.
+            kernel_size (tuple): Size of the convolutional kernel.
+            x (Tensor): Input tensor.
+
+        Returns:
+            list: List of convolutional stack layers.
+        """
         # Stride=1, padding=(1,1)
         conv = Conv2dLayer(in_channels, out_filters, kernel_size, 1, (1, 1), x)
         relu = ReLULayer(conv)

poet/poet_solver.py

@@ -107,7 +107,7 @@ def _initialize_variables(self):
 
     def _create_correctness_constraints(self):
         # ensure all computations are possible
-        for (u, v) in self.g.edge_list:
+        for u, v in self.g.edge_list:
             for t in range(self.T):
                 self.m += self.R[t][v] <= self.R[t][u] + self.SRam[t][u]
         # ensure all checkpoints are in memory

poet/poet_solver_gurobi.py

@@ -13,6 +13,7 @@
 
 # noinspection PyPackageRequirements
 
+
 # POET ILP defined using Gurobi
 class POETSolverGurobi:
     def __init__(
@@ -124,7 +125,7 @@ def _disable_paging(self):
 
     def _create_correctness_constraints(self):
         # ensure all computations are possible
-        for (u, v) in self.g.edge_list:
+        for u, v in self.g.edge_list:
             for t in range(self.T):
                 self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL, self.R[t, u] + self.SRam[t, u])
         # ensure all checkpoints are in memory

poet/utils/checkmate/core/utils/graph.py

@@ -10,7 +10,7 @@
 def edge_to_adj_list(E: EdgeList, convert_undirected=False):
     """Returns an (undirected / bidirectional) adjacency list"""
     adj_list = defaultdict(set)
-    for (i, j) in E:
+    for i, j in E:
         adj_list[i].add(j)
         if convert_undirected:
             adj_list[j].add(i)

poet/utils/checkmate/core/utils/timer.py

@@ -4,7 +4,22 @@
 
 
 class Timer:
+    """
+    A context manager class for measuring and printing elapsed time durations.
+
+    Attributes:
+        niters (int): Number of iterations.
+        _elapsed (Decimal): Total elapsed time.
+        _name (str): Name of the timer.
+        _print_results (bool): Whether to print results after timing.
+        _start_time (float): Start time of the timer.
+        _children (dict): Dictionary to hold child timers.
+        _count (int): Number of times the timer has been started.
+    """
     def __init__(self, name, extra_data=None, print_results=False, niters=None):
+        """
+        Initialize a Timer object.
+        """
         self.niters = niters
         self._elapsed = Decimal()
         self._name = name
@@ -17,18 +32,27 @@ def __init__(self, name, extra_data=None, print_results=False, niters=None):
 
     @property
     def elapsed(self):
+        """Return the total elapsed time in seconds."""
         return float(self._elapsed)
 
     def __enter__(self):
+        """Enter the context manager and start the timer."""
         self.start()
         return self
 
     def __exit__(self, *_):
+        """Exit the context manager and stop the timer."""
         self.stop()
         if self._print_results:
             self.print_results()
 
     def child(self, name):
+        """
+        Create a child Timer with a given name.
+
+        Returns:
+            Timer: Child Timer instance.
+        """
         try:
             return self._children[name]
         except KeyError:
@@ -47,6 +71,12 @@ def print_results(self):
         print(self._format_results())
 
     def _format_results(self, indent="  "):
+        """
+        Format the timer results for printing.
+
+        Returns:
+            str: Formatted timer results.
+        """
         children = self._children.values()
         elapsed = self._elapsed or sum(c._elapsed for c in children)
         result = "%s: %.3fs" % (self._name, elapsed)