Migrated type and metadata field to NIRNode dataclass field factory

nir/ir/conv.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass, field
-from typing import Any, Dict, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
 
 import numpy as np
 
@@ -41,9 +41,6 @@ class Conv1d(NIRNode):
     dilation: int  # Dilation
     groups: int  # Groups
     bias: np.ndarray  # Bias C_out
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         if isinstance(self.padding, str) and self.padding not in ["same", "valid"]:
@@ -98,15 +95,13 @@ class Conv2d(NIRNode):
     :type bias: np.ndarray
     """
 
-    # Shape of input tensor (overrrides input_type from
     input_shape: Optional[Tuple[int, int]]  # N_x, N_y
     weight: np.ndarray  # Weight C_out * C_in * W_x * W_y
     stride: Union[int, Tuple[int, int]]  # Stride
     padding: Union[int, Tuple[int, int], str]  # Padding
     dilation: Union[int, Tuple[int, int]]  # Dilation
     groups: int  # Groups
     bias: np.ndarray  # Bias C_out
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         if isinstance(self.padding, str) and self.padding not in ["same", "valid"]:

nir/ir/delay.py

@@ -1,5 +1,4 @@
-from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -17,9 +16,6 @@ class Delay(NIRNode):
     """
 
     delay: np.ndarray  # Delay
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         # set input and output shape, if not set by user

nir/ir/flatten.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from dataclasses import dataclass
+from typing import Any, Dict
 
 import numpy as np
 
@@ -16,14 +16,11 @@ class Flatten(NIRNode):
     input_type must be a dict with one key: "input".
     """
 
-    # Shape of input tensor (overrrides input_type from
-    # NIRNode to allow for non-keyword (positional) initialization)
+    # Shape of input tensor. Overrides the keyword-only input_type from NIRNode
+    # so it can be passed positionally
     input_type: Types
     start_dim: int = 1  # First dimension to flatten
     end_dim: int = -1  # Last dimension to flatten
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.input_type = parse_shape_argument(self.input_type, "input")

nir/ir/graph.py

@@ -1,5 +1,5 @@
 from collections import Counter
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Any, Dict, Optional
 
 import numpy as np
@@ -34,9 +34,6 @@ class NIRGraph(NIRNode):
 
     nodes: Nodes  # List of computational nodes
     edges: Edges  # List of edges between nodes
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __init__(
         self,
@@ -501,10 +498,9 @@ class Input(NIRNode):
     This is a virtual node, which allows feeding in data into the graph.
     """
 
-    # Shape of incoming data (overrrides input_type from
-    # NIRNode to allow for non-keyword (positional) initialization)
+    # Shape of incoming data. Overrides the keyword-only input_type from
+    # NIRNode so it can be passed positionally as the node's defining argument.
     input_type: Types
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.input_type = parse_shape_argument(self.input_type, "input")
@@ -529,10 +525,9 @@ class Output(NIRNode):
     Defines an output of the graph.
     """
 
-    # Type of incoming data (overrrides input_type from
-    # NIRNode to allow for non-keyword (positional) initialization)
+    # Type of outgoing data. Overrides the keyword-only output_type from
+    # NIRNode so it can be passed positionally
     output_type: Types
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.output_type = parse_shape_argument(self.output_type, "output")
@@ -557,6 +552,8 @@ class Identity(NIRNode):
     This is a virtual node, which allows for the identity operation.
     """
 
+    # Shape of incoming data. Overrides the keyword-only input_type from
+    # NIRNode so it can be passed positionally
     input_type: Types
 
     def __post_init__(self):

nir/ir/linear.py

@@ -1,5 +1,4 @@
-from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -21,9 +20,6 @@ class Affine(NIRNode):
 
     weight: np.ndarray  # Weight term
     bias: np.ndarray  # Bias term
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         assert len(self.weight.shape) >= 2, "Weight must be at least 2D"
@@ -46,7 +42,6 @@ class Linear(NIRNode):
     """
 
     weight: np.ndarray  # Weight term
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         assert len(self.weight.shape) >= 2, "Weight must be at least 2D"
@@ -70,7 +65,6 @@ class Scale(NIRNode):
     """
 
     scale: np.ndarray  # Scaling factor
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.input_type = {"input": np.array(self.scale.shape)}

nir/ir/neuron.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Any, Dict, Optional
 
 import numpy as np
@@ -32,9 +32,6 @@ class CubaLI(NIRNode):
     r: np.ndarray  # Resistance
     v_leak: np.ndarray  # Leak voltage
     w_in: np.ndarray = 1.0  # Input current weight
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         assert (
@@ -91,9 +88,6 @@ class CubaLIF(NIRNode):
     v_threshold: np.ndarray  # Firing threshold
     v_reset: Optional[np.ndarray] = None  # Reset potential
     w_in: np.ndarray = 1.0  # Input current weight
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         if self.v_reset is None:
@@ -129,9 +123,6 @@ class I(NIRNode):  # noqa: E742
     """
 
     r: np.ndarray
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.input_type = {"input": np.array(self.r.shape)}
@@ -163,9 +154,6 @@ class IF(NIRNode):
     r: np.ndarray  # Resistance
     v_threshold: np.ndarray  # Firing threshold
     v_reset: Optional[np.ndarray] = None  # Reset potential
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         if self.v_reset is None:
@@ -200,9 +188,6 @@ class LI(NIRNode):
     tau: np.ndarray  # Time constant
     r: np.ndarray  # Resistance
     v_leak: np.ndarray  # Leak voltage
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         assert (
@@ -247,9 +232,6 @@ class LIF(NIRNode):
     v_leak: np.ndarray  # Leak voltage
     v_threshold: np.ndarray  # Firing threshold
     v_reset: Optional[np.ndarray] = None  # Reset potential
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         if self.v_reset is None:

nir/ir/node.py

@@ -1,7 +1,9 @@
 from abc import ABC
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from typing import Any, Dict
 
+import numpy as np
+
 
 @dataclass(eq=False)
 class NIRNode(ABC):
@@ -11,11 +13,14 @@ class NIRNode(ABC):
     instantiated.
     """
 
-    # Note: Adding input/output types and metadata as follows is ideal, but requires Python 3.10
-    # TODO: implement this in 2025 when 3.9 is EOL
-    # input_type: Dict[str, np.ndarray] = field(init=False, kw_only=True)
-    # output_type: Dict[str, np.ndarray] = field(init=False, kw_only=True)
-    # metadata: Dict[str, Any] = field(init=True, default_factory=dict)
+    # input_type and output_type are not initialized in the constructor, they are
+    # populated in __post_init__ for each subclass. metadata is an optional
+    # keyword argument. All three are keyword-only so that subclasses can add
+    # positional fields without running into the "non-default argument follows
+    # default argument" ordering error. (Requires Python 3.10+.)
+    input_type: Dict[str, np.ndarray] = field(init=False, kw_only=True)
+    output_type: Dict[str, np.ndarray] = field(init=False, kw_only=True)
+    metadata: Dict[str, Any] = field(default_factory=dict, kw_only=True)
 
     def __init__(self) -> None:
         raise AttributeError("NIRNode does not have a default constructor.")

nir/ir/pooling.py

@@ -1,5 +1,4 @@
-from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -13,9 +12,6 @@ class SumPool2d(NIRNode):
     kernel_size: np.ndarray  # (Height, Width)
     stride: np.ndarray  # (Height, width)
     padding: np.ndarray  # (Height, width)
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.input_type = {"input": None}
@@ -29,7 +25,6 @@ class AvgPool2d(NIRNode):
     kernel_size: np.ndarray  # (Height, Width)
     stride: np.ndarray  # (Height, width)
     padding: np.ndarray  # (Height, width)
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.input_type = {"input": None}

nir/ir/threshold.py

@@ -1,5 +1,4 @@
-from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -20,9 +19,6 @@ class Threshold(NIRNode):
     """
 
     threshold: np.ndarray  # Firing threshold
-    input_type: Optional[Dict[str, np.ndarray]] = None
-    output_type: Optional[Dict[str, np.ndarray]] = None
-    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         self.input_type = {"input": np.array(self.threshold.shape)}

pyproject.toml

@@ -32,7 +32,7 @@ classifiers = [
 ]
 dependencies = [ "numpy", "h5py" ]
 dynamic = ["version"] # Version number read from __init__.py
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 
 [project.urls]
 homepage = "https://github.com/neuromorphs/nir"

tests/test_ir.py

@@ -14,6 +14,27 @@ def test_has_NIRNode():
     assert hasattr(nir, "NIRNode")
 
 
+def test_input_output_type_mutable():
+    """The inherited input_type/output_type fields are populated in __post_init__
+    and must remain writable from the outside (graph type inference reassigns them
+    and mutates their keys in place)."""
+    node = mock_affine(2, 3)
+
+    # Mutate an existing key in place
+    node.input_type["input"] = np.array([42])
+    assert np.array_equal(node.input_type["input"], np.array([42]))
+
+    # Add a new key
+    node.output_type["extra"] = np.array([7])
+    assert np.array_equal(node.output_type["extra"], np.array([7]))
+
+    # Reassign the whole attribute
+    node.input_type = {"input": np.array([1, 1])}
+    node.output_type = {"output": np.array([1, 1])}
+    assert np.array_equal(node.input_type["input"], np.array([1, 1]))
+    assert np.array_equal(node.output_type["output"], np.array([1, 1]))
+
+
 def test_eq():
     a = nir.Input(np.array([2, 3]))
     a2 = nir.Input(np.array([2, 3]))