Renamed surrogate gradient. Added Identity (#114)

nir/ir/__init__.py

@@ -7,7 +7,7 @@
 from .linear import Affine, Linear, Scale
 from .neuron import IF, LI, LIF, CubaLIF, I
 from .pooling import AvgPool2d, SumPool2d
-from .surrogate_gradient import Threshold
+from .threshold import Threshold
 from .typing import NIRNode
 
 # NIRNodes that can be (de)serialized

nir/ir/graph.py

@@ -496,3 +496,23 @@ def from_dict(cls, node: Dict[str, Any]) -> "NIRNode":
         node["output_type"] = {"output": node["shape"]}
         del node["shape"]
         return super().from_dict(node)
+
+@dataclass(eq=False)
+class Identity(NIRNode):
+    """Identity Node.
+
+    This is a virtual node, which allows for the identity operation.
+    """
+    input_type: Types
+
+    def __post_init__(self):
+        self.output_type = self.input_type
+
+    def to_dict(self) -> Dict[str, Any]:
+        ret = super().to_dict()
+        ret["shape"] = self.output_type["output"]
+        return ret
+
+    @classmethod
+    def from_dict(cls, node: Dict[str, Any]) -> "NIRNode":
+        return super().from_dict(node)