[CustomOp] Adding comments to im2col

src/finn/custom_op/im2col.py

@@ -11,13 +11,14 @@ from finn.core.datatype import DataType
 
 
 def compute_conv_output_dim(ifm_dim, k, stride, pad=0):
-    "Return spatial output dimension size for convolution with given params."
+    """Returns spatial output dimension size for convolution with given params."""
     return int(((ifm_dim + 2 * pad - k) / stride) + 1)
 
 
 def get_im2col_indices_nchw(
     x_shape, field_height, field_width, padding=0, stride_y=1, stride_x=1
 ):
+    """Returns im2col indices."""
     # First figure out what the size of the output should be
     N, C, H, W = x_shape
     assert (H + 2 * padding - field_height) % stride_y == 0
@@ -41,7 +42,9 @@ def get_im2col_indices_nchw(
 def im2col_indices_nchw(
     x, field_height, field_width, padding=0, stride_y=1, stride_x=1, pad_val=0
 ):
-    """ An implementation of im2col based on some fancy indexing """
+    """Performs im2col on x with given field height and width, as well as values
+    for padding and stride size.
+    Returns result of im2col."""
     # Zero-pad the input
     p = padding
     x_padded = np.pad(