Add AvgPool2d

docs/source/primitives.md

@@ -18,6 +18,7 @@ NIR defines 16 fundamental primitives listed in the table below, which backends
 | **Leaky integrate-fire (LIF)** | $\tau, \text{R}, v_\text{leak}, v_\text{thr}$ | **LI**; **Threshold** | $\begin{cases} v-v_\text{thr} & \text{Spike} \\ v & \text{else} \end{cases}$ |
 | **Scale** | $s$ | $s I$ | - |
 | **SumPooling** | $p$ | $\sum_{j} x_j$ |  |
+| **AvgPooling** | $p$ | **SumPooling**; **Scale** | - |
 | **Threshold** | $\theta_\text{thr}$ | $H(I - \theta_\text{thr})$ | - |
 
 Each primitive is defined by their own dynamical equation, specified in the [API docs](https://nnir.readthedocs.io/en/latest/modindex.html).
@@ -33,4 +34,4 @@ $$
 $$
 
 ## Format
-The intermediate represenation can be stored as hdf5 file, which benefits from compression. 
+The intermediate represenation can be stored as hdf5 file, which benefits from compression. 

nir/ir/__init__.py

@@ -6,7 +6,7 @@
 from .graph import Input, NIRGraph, Output
 from .linear import Affine, Linear, Scale
 from .neuron import IF, LI, LIF, CubaLIF, I
-from .pooling import SumPool2d
+from .pooling import AvgPool2d, SumPool2d
 from .surrogate_gradient import Threshold
 from .typing import NIRNode
 
@@ -34,6 +34,7 @@
     "LI",
     "LIF",
     # pooling
+    "AvgPool2d",
     "SumPool2d",
     # surrogate_gradient
     "Threshold",
@@ -81,6 +82,7 @@ def dict2NIRNode(data_dict: Dict[str, Any]) -> NIRNode:
     "LI",
     "LIF",
     # pooling
+    "AvgPool2d",
     "SumPool2d",
     # surrogate_gradient
     "Threshold",

nir/ir/graph.py

@@ -7,7 +7,7 @@
 from .conv import Conv1d, Conv2d
 from .flatten import Flatten
 from .node import NIRNode
-from .pooling import SumPool2d
+from .pooling import AvgPool2d, SumPool2d
 from .typing import Edges, Nodes, Types
 from .utils import (
     calc_flatten_output,
@@ -407,6 +407,19 @@ def _forward_type_inference(self, debug=True):
                     )
                     post_node.output_type = {"output": output_type}
 
+                elif isinstance(post_node, AvgPool2d):
+                    output_shape = calculate_conv_output(
+                        pre_node.output_type["output"][1:],
+                        post_node.padding,
+                        1,
+                        post_node.kernel_size,
+                        post_node.stride,
+                    )
+                    output_type = np.array(
+                        [post_node.input_type["input"][0], *output_shape]
+                    )
+                    post_node.output_type = {"output": output_type}
+
                 elif isinstance(post_node, Flatten):
                     print("updateing flatten output")
                     post_node.output_type = {

nir/ir/pooling.py

@@ -16,3 +16,16 @@ class SumPool2d(NIRNode):
     def __post_init__(self):
         self.input_type = {"input": None}
         self.output_type = {"output": None}
+
+
+@dataclass(eq=False)
+class AvgPool2d(NIRNode):
+    """Average pooling layer in 2d."""
+
+    kernel_size: np.ndarray  # (Height, Width)
+    stride: np.ndarray  # (Height, width)
+    padding: np.ndarray  # (Height, width)
+
+    def __post_init__(self):
+        self.input_type = {"input": None}
+        self.output_type = {"output": None}

nir/serialization.py

@@ -45,6 +45,12 @@ def read_node(node: Any) -> nir.typing.NIRNode:
             stride=node["stride"][()],
             padding=node["padding"][()],
         )
+    elif node["type"][()] == b"AvgPool2d":
+        return nir.AvgPool2d(
+            kernel_size=node["kernel_size"][()],
+            stride=node["stride"][()],
+            padding=node["padding"][()],
+        )
     elif node["type"][()] == b"Delay":
         return nir.Delay(delay=node["delay"][()])
     elif node["type"][()] == b"Flatten":