Fix: Scalar Parameter Broadcasting Forces Inefficient Weight Duplication

[siddiquifaras]

Summary

NIR's type inference system currently requires neuron parameters (tau, v_threshold, etc.) to be arrays with explicit shapes to determine the number of neurons in a layer. This forces scalar parameters to be broadcast into full arrays, creating unnecessary weight duplication that is inefficient for hardware deployment.

Problem Description

When exporting spiking neural networks to NIR, scalar parameters (which are common and hardware-efficient) must be converted to full arrays to satisfy NIR's type checking system.

Example

Original snnTorch model:

import snntorch as snn
lif = snn.Leaky(beta=0.9)  # Single scalar parameter for all neurons

Current NIR export requirement:

# For a 128-neuron layer, scalar beta=0.9 becomes:
nir.LIF(
    tau=np.array([0.001, 0.001, 0.001, ..., 0.001]),  # 128 identical values
    v_threshold=np.array([1.0, 1.0, 1.0, ..., 1.0]),  # 128 identical values
    # ... all parameters duplicated 128 times
)

What we want:

nir.LIF(
    tau=np.array(0.001),        # Single scalar value
    v_threshold=np.array(1.0),  # Single scalar value
    neuron_count=128,           # Explicit neuron count
)

Root Cause

NIR's type inference relies on parameter array shapes to determine neuron counts:

# From NIR source - infers neuron count from parameter shapes
def infer_neuron_count(self):
    return self.tau.shape[0] if self.tau.ndim > 0 else None

When parameters are scalars (0-d arrays), NIR cannot infer the neuron count and sets empty type annotations:

• input_type: {'input': array([], dtype=float64)} (wrong)
• output_type: {'output': array([], dtype=float64)} (wrong)

This causes type checking to fail with errors like:

ValueError: Type inference error: type mismatch: fc1.output: [[128]] -> lif.input: []

Impact

1. Memory Inefficiency

• A scalar parameter becomes an array of N identical values
• For large layers (e.g., 1024 neurons), this creates 1024x memory overhead
• Particularly problematic for neuromorphic hardware with limited memory

2. Hardware Deployment Issues

• Neuromorphic chips are optimized for shared scalar parameters
• Broadcasting forces hardware to store/access redundant identical weights
• Reduces effective model capacity on memory-constrained devices

3. Export Complexity

• Forces export tools to choose between efficiency and type safety
• Current workarounds require parameter broadcasting that loses semantic meaning

Proposed Solutions

Option 1: Explicit Neuron Count Parameter

Add an optional neuron_count parameter to neuron nodes:

nir.LIF(
    tau=np.array(0.001),     # Scalar - applies to all neurons
    v_threshold=np.array(1.0), 
    neuron_count=128,        # Explicit count for type inference
    input_type={'input': np.array([128])},   # Auto-generated from neuron_count
    output_type={'output': np.array([128])}, # Auto-generated from neuron_count
)

Option 2: Broadcast Semantics

Add support for broadcast-compatible parameters:

nir.LIF(
    tau=np.array(0.001),           # Scalar broadcasts to all neurons
    v_threshold=np.array([1.0, 1.2, 1.0]),  # Per-neuron values where needed
    broadcast_params=['tau'],       # Explicit broadcast annotation
)

Option 3: Shape Inference Context

Allow type inference to use context from connected layers:

# Infer neuron count from connected Affine layer output shape
fc1 = nir.Affine(weight=np.random.randn(128, 784))  # 128 outputs
lif = nir.LIF(tau=np.array(0.001))  # Infer 128 neurons from fc1 connection

Current Workaround

Export tools currently broadcast scalars to satisfy type checking:

def to_array(val, n_neurons):
    """Broadcast scalar to array of n_neurons length"""
    if np.isscalar(val):
        return np.full(n_neurons, float(val))
    # ... handle existing arrays

This works but loses the semantic meaning and efficiency of scalar parameters.

Test Case

import nir
import numpy as np

# This should work without broadcasting
lif_scalar = nir.LIF(
    tau=np.array(0.001),        # Scalar parameter
    v_threshold=np.array(1.0),
    neuron_count=128,           # Proposed: explicit neuron count
)

# Type inference should succeed
graph = nir.NIRGraph(
    nodes={'lif': lif_scalar, ...},
    edges=[...]
)
# Should pass: nir.read(file, type_check=True)

References

• snnTorch export issue: Scalar parameters must be broadcast to arrays for NIR compatibility
• Hardware efficiency: Neuromorphic chips (Intel Loihi, SpiNNaker) optimize for shared scalar parameters
• Memory impact: Large models with broadcast parameters can exceed hardware memory limits

Environment

• NIR version: 1.0.7
• nirtorch version: 2.0.5
• Use case: snnTorch → NIR → neuromorphic hardware deployment

---

Expected Outcome: NIR should support efficient scalar parameters while maintaining type safety, enabling optimal hardware deployment without unnecessary memory overhead.

[Jegp]

Thank you for bringing this to our attention! I can see why it might seem backwards to have this constraint. It's put in place because we wanted to disambiguate 100%. But, we're now working towards clearer type distinctions, which will state clearly how many inputs/outputs there are, so this should, again, make the scalar broadcasting safe.

One thing I would like to say, though, is that I don't think NIR is responsible for hardware efficiency. It's a representation between platforms, so NIR has no opinion of influence in what's put on hardware. That's entirely on Loihi's side (although I can see how it might be hard to "de-broadcast"). NIR should, of course, not be monumentally huge. So, say we implement this, do you know how many bits we save? Is it worth the effort.

That said, I'm not a huge fan of the first two options you're suggestion. Option 1 seems verbose and the shape argument should really be a type or a piece of meta-data rather than something the user should care about. Option 2 feels wrong for the same reason: the user shouldn't have to do this for every single parameter.

My favorite approach is 3: we infer the right shape from the arguments. But, we have to be careful. In your example code, how do we know that there are 128 neurons? The fc1 and lif nodes aren't connected. One idea would be to hook into the NIRGraph's infer_types method that is triggered when someone constructs a graph (unless the type_check flag is manually disabled). This actually makes sense to me: the type mostly becomes a problem when someone tries to put together two nodes because they need to have the same I/O pattern. This is also where your error is generated from: we "just" need to fix the typing to go through all the connections in the graph and add a clause that tests for either full shape information OR (scalar + information to infer data).

One small caveat we should be aware of in that case, is when we try to construct a simple network like, say Conv2d -> LIF. Conv2d doesn't in itself have any output type (we just know the size of the kernels, but not the size of the inputs - PyTorch does the same). So, what then? I think the best idea would be to specify the type using the output_type argument or so, but it's not pretty...

[siddiquifaras]

Hey Jens, thanks for the follow-up!

I understand the design rationale and agree that NIR should remain agnostic to hardware efficiency. The backend is responsible for de-broadcasting.

On Memory Savings

You asked whether it is worth the effort. For a 1024-neuron layer with 5 float32 parameters:

• Broadcast: 1024 x 5 x 4 = 20,480 bytes
• Scalar: 5 x 4 = 20 bytes

For a 10-layer network: 200KB vs 200 bytes :)
The savings are significant for edge deployment, but I acknowledge this is a backend concern, not NIR's responsibility.

Support for Option 3

I agree Option 3 would be the right approach given your thoughts. Below is a production-ready implementation proposal.

Design Principles

1. Backward compatible: Existing code with explicit array parameters continues to work unchanged
2. Fail-safe: Type inference fails gracefully with clear error messages when context is insufficient
3. Deterministic: Same graph always produces same inferred types regardless of iteration order
4. Extensible: New node types can implement shape inference by following a simple protocol

Architecture

NIRGraph.infer_types()
    |
    +-- TopologicalIterator (handles cycles via SCCs)
    |
    +-- For each node in order:
            |
            +-- node.resolve_input_type(predecessor_outputs)
            +-- node.compute_output_type()
            +-- Validate or raise InferenceError

Implementation

1. Base Protocol for Shape Inference

# nir/ir/typing.py

from abc import abstractmethod
from typing import Dict, List, Optional, Protocol
import numpy as np

class ShapeArray(np.ndarray):
    """Typed array for shape information with broadcasting semantics."""
    
    @classmethod
    def from_value(cls, val) -> "ShapeArray":
        arr = np.atleast_1d(np.asarray(val))
        return arr.view(cls)
    
    @property
    def is_scalar(self) -> bool:
        return self.size == 1
    
    def broadcast_to(self, shape: tuple) -> "ShapeArray":
        if self.is_scalar:
            return np.broadcast_to(self, shape).view(ShapeArray)
        if self.shape == shape:
            return self
        raise ValueError(f"Cannot broadcast {self.shape} to {shape}")


class InferenceError(Exception):
    """Raised when type inference fails due to insufficient context."""
    def __init__(self, node_name: str, reason: str):
        self.node_name = node_name
        self.reason = reason
        super().__init__(f"Type inference failed for '{node_name}': {reason}")


class ShapeInferable(Protocol):
    """Protocol for nodes that support shape inference from context."""
    
    input_type: Optional[Dict[str, np.ndarray]]
    output_type: Optional[Dict[str, np.ndarray]]
    
    def resolve_input_type(self, predecessor_outputs: List[Dict[str, np.ndarray]]) -> None:
        """Resolve input_type from predecessor output types."""
        ...
    
    def compute_output_type(self) -> None:
        """Compute output_type from input_type and node parameters."""
        ...
    
    def validate_types(self) -> None:
        """Validate that input_type and output_type are fully resolved."""
        ...

2. Graph Traversal with Cycle Handling

# nir/ir/graph.py

from typing import List, Set, Dict, Tuple
from collections import defaultdict

class NIRGraph:
    
    def infer_types(self) -> None:
        """Infer input/output types for all nodes using graph context.
        
        Traverses nodes in topological order, propagating shape information
        forward. Handles cyclic graphs by processing strongly connected
        components and requiring at least one node per cycle to have
        explicit types.
        
        Raises:
            InferenceError: If type inference fails for any node.
        """
        adjacency = self._build_adjacency()
        order = self._topological_order_with_cycles(adjacency)
        
        for node_name in order:
            self._infer_node_types(node_name, adjacency)
        
        self._validate_all_edges()
    
    def _build_adjacency(self) -> Dict[str, Tuple[List[str], List[str]]]:
        """Build predecessor/successor lists for each node."""
        predecessors = defaultdict(list)
        successors = defaultdict(list)
        
        for src, dst in self.edges:
            successors[src].append(dst)
            predecessors[dst].append(src)
        
        return {
            name: (predecessors[name], successors[name])
            for name in self.nodes
        }
    
    def _topological_order_with_cycles(
        self, 
        adjacency: Dict[str, Tuple[List[str], List[str]]]
    ) -> List[str]:
        """Compute processing order using Tarjan's SCC algorithm.
        
        Returns nodes in an order where:
        1. Non-cyclic nodes come after their predecessors
        2. Cyclic nodes (SCCs) are grouped together
        """
        index_counter = [0]
        stack = []
        lowlinks = {}
        index = {}
        on_stack = set()
        sccs = []
        
        def strongconnect(node: str):
            index[node] = index_counter[0]
            lowlinks[node] = index_counter[0]
            index_counter[0] += 1
            stack.append(node)
            on_stack.add(node)
            
            predecessors, successors = adjacency[node]
            for successor in successors:
                if successor not in index:
                    strongconnect(successor)
                    lowlinks[node] = min(lowlinks[node], lowlinks[successor])
                elif successor in on_stack:
                    lowlinks[node] = min(lowlinks[node], index[successor])
            
            if lowlinks[node] == index[node]:
                scc = []
                while True:
                    w = stack.pop()
                    on_stack.remove(w)
                    scc.append(w)
                    if w == node:
                        break
                sccs.append(scc)
        
        for node in self.nodes:
            if node not in index:
                strongconnect(node)
        
        # Flatten SCCs in reverse order (respects dependencies)
        result = []
        for scc in reversed(sccs):
            result.extend(scc)
        
        return result
    
    def _infer_node_types(
        self, 
        node_name: str, 
        adjacency: Dict[str, Tuple[List[str], List[str]]]
    ) -> None:
        """Infer types for a single node from its predecessors."""
        node = self.nodes[node_name]
        predecessors, _ = adjacency[node_name]
        
        # Collect predecessor output types
        pred_outputs = []
        for pred_name in predecessors:
            pred_node = self.nodes[pred_name]
            if pred_node.output_type is not None:
                pred_outputs.append(pred_node.output_type)
        
        # Resolve input type from predecessors
        if hasattr(node, 'resolve_input_type'):
            node.resolve_input_type(pred_outputs)
        elif node.input_type is None and pred_outputs:
            # Default: take first predecessor's output type
            node.input_type = {
                "input": next(iter(pred_outputs[0].values())).copy()
            }
        
        # Compute output type from input type
        if hasattr(node, 'compute_output_type'):
            node.compute_output_type()
        
        # Validate
        if hasattr(node, 'validate_types'):
            node.validate_types()
    
    def _validate_all_edges(self) -> None:
        """Validate type compatibility across all edges."""
        for src, dst in self.edges:
            src_node = self.nodes[src]
            dst_node = self.nodes[dst]
            
            if src_node.output_type is None:
                raise InferenceError(src, "output_type not resolved")
            if dst_node.input_type is None:
                raise InferenceError(dst, "input_type not resolved")
            
            src_shape = next(iter(src_node.output_type.values()))
            dst_shape = next(iter(dst_node.input_type.values()))
            
            if not self._shapes_compatible(src_shape, dst_shape):
                raise InferenceError(
                    dst,
                    f"Shape mismatch: {src}.output {src_shape} incompatible with {dst}.input {dst_shape}"
                )
    
    @staticmethod
    def _shapes_compatible(a: np.ndarray, b: np.ndarray) -> bool:
        """Check if shapes are compatible under broadcasting rules."""
        a = np.atleast_1d(a)
        b = np.atleast_1d(b)
        
        # Scalar is compatible with anything
        if a.size == 1 or b.size == 1:
            return True
        
        # Same shape
        if a.shape == b.shape and np.array_equal(a, b):
            return True
        
        return False

3. Neuron Node Implementation

# nir/ir/neuron.py

from dataclasses import dataclass
from typing import Dict, List, Optional
import numpy as np

from .node import NIRNode
from .typing import InferenceError, ShapeArray

@dataclass(eq=False)
class LIF(NIRNode):
    """Leaky Integrate-and-Fire neuron.
    
    Parameters can be scalars (shared across all neurons) or arrays
    (per-neuron values). Shape is inferred from:
    1. First non-scalar parameter, or
    2. Predecessor output type during graph inference
    """
    tau: np.ndarray
    r: np.ndarray
    v_leak: np.ndarray
    v_threshold: np.ndarray
    
    def __post_init__(self):
        # Convert to arrays
        self.tau = np.atleast_1d(np.asarray(self.tau, dtype=np.float64))
        self.r = np.atleast_1d(np.asarray(self.r, dtype=np.float64))
        self.v_leak = np.atleast_1d(np.asarray(self.v_leak, dtype=np.float64))
        self.v_threshold = np.atleast_1d(np.asarray(self.v_threshold, dtype=np.float64))
        
        # Try to infer shape from parameters
        shape = self._infer_shape_from_params()
        if shape is not None:
            self.input_type = {"input": np.array(shape)}
            self.output_type = {"output": np.array(shape)}
        # Otherwise, leave as None for graph-level inference
    
    def _infer_shape_from_params(self) -> Optional[tuple]:
        """Extract shape from first non-scalar parameter."""
        for param in [self.tau, self.r, self.v_leak, self.v_threshold]:
            if param.size > 1:
                return param.shape
        return None
    
    def resolve_input_type(self, predecessor_outputs: List[Dict[str, np.ndarray]]) -> None:
        """Resolve input_type from predecessor outputs if not already set."""
        if self.input_type is not None:
            return
        
        if not predecessor_outputs:
            return
        
        # Use first predecessor's output shape
        pred_shape = next(iter(predecessor_outputs[0].values()))
        self.input_type = {"input": np.asarray(pred_shape).copy()}
    
    def compute_output_type(self) -> None:
        """LIF output shape equals input shape."""
        if self.output_type is not None:
            return
        
        if self.input_type is not None:
            self.output_type = {"output": self.input_type["input"].copy()}
    
    def validate_types(self) -> None:
        """Validate that types are resolved and parameters are compatible."""
        if self.input_type is None:
            raise InferenceError("LIF", "input_type not resolved and no non-scalar parameters")
        
        expected_size = self.input_type["input"][-1] if len(self.input_type["input"]) > 0 else 1
        
        for name, param in [
            ("tau", self.tau), 
            ("r", self.r), 
            ("v_leak", self.v_leak), 
            ("v_threshold", self.v_threshold)
        ]:
            if param.size != 1 and param.shape[-1] != expected_size:
                raise InferenceError(
                    "LIF", 
                    f"Parameter '{name}' shape {param.shape} incompatible with neuron count {expected_size}"
                )

4. Conv2d Node Implementation

# nir/ir/conv.py

from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import numpy as np

from .node import NIRNode
from .typing import InferenceError

@dataclass(eq=False)
class Conv2d(NIRNode):
    """2D Convolution layer.
    
    Output shape is computed from input shape and convolution parameters.
    Requires input_type to be set (either explicitly or from predecessors).
    """
    weight: np.ndarray  # [out_channels, in_channels/groups, kH, kW]
    bias: np.ndarray    # [out_channels]
    stride: Tuple[int, int] = (1, 1)
    padding: Tuple[int, int] = (0, 0)
    dilation: Tuple[int, int] = (1, 1)
    groups: int = 1
    
    def __post_init__(self):
        self.weight = np.asarray(self.weight)
        self.bias = np.asarray(self.bias)
        
        if self.weight.ndim != 4:
            raise ValueError(f"Conv2d weight must be 4D, got {self.weight.ndim}D")
        
        # Cannot compute output_type without input_type
        # Will be resolved during graph inference
    
    def resolve_input_type(self, predecessor_outputs: List[Dict[str, np.ndarray]]) -> None:
        """Resolve input_type from predecessor outputs."""
        if self.input_type is not None:
            return
        
        if not predecessor_outputs:
            return
        
        pred_shape = next(iter(predecessor_outputs[0].values()))
        self.input_type = {"input": np.asarray(pred_shape).copy()}
    
    def compute_output_type(self) -> None:
        """Compute output shape from input shape and conv parameters."""
        if self.output_type is not None:
            return
        
        if self.input_type is None:
            return
        
        input_shape = self.input_type["input"]
        
        # Expected input: [C, H, W] or [H, W]
        if len(input_shape) == 3:
            _, h_in, w_in = input_shape
        elif len(input_shape) == 2:
            h_in, w_in = input_shape
        else:
            raise InferenceError(
                "Conv2d", 
                f"Expected 2D or 3D input shape, got {input_shape}"
            )
        
        out_channels, _, kH, kW = self.weight.shape
        
        h_out = (h_in + 2 * self.padding[0] - self.dilation[0] * (kH - 1) - 1) // self.stride[0] + 1
        w_out = (w_in + 2 * self.padding[1] - self.dilation[1] * (kW - 1) - 1) // self.stride[1] + 1
        
        self.output_type = {"output": np.array([out_channels, h_out, w_out])}
    
    def validate_types(self) -> None:
        """Validate input/output types are resolved."""
        if self.input_type is None:
            raise InferenceError(
                "Conv2d", 
                "input_type not resolved. Conv2d requires explicit input_type or a predecessor with output_type."
            )
        if self.output_type is None:
            raise InferenceError("Conv2d", "output_type not computed")

5. Integration with Existing API

# nir/ir/graph.py (continued)

class NIRGraph:
    
    def __post_init__(self):
        # ... existing initialization ...
        
        if self.type_check:
            self.infer_types()
            # Existing edge validation now covered by _validate_all_edges()

Usage Examples

Scalar parameters with graph context:

import nir
import numpy as np

# Scalar parameters are valid when graph provides context
nodes = {
    "input": nir.Input(input_type={"input": np.array([784])}),
    "fc1": nir.Affine(
        weight=np.random.randn(128, 784).astype(np.float32),
        bias=np.zeros(128, dtype=np.float32)
    ),
    "lif1": nir.LIF(
        tau=np.array(0.01),       # Scalar - will infer 128 neurons from fc1
        r=np.array(1.0),
        v_leak=np.array(0.0),
        v_threshold=np.array(1.0)
    ),
    "output": nir.Output(output_type={"output": np.array([128])})
}

edges = [("input", "fc1"), ("fc1", "lif1"), ("lif1", "output")]

# Type inference happens automatically
graph = nir.NIRGraph(nodes, edges, type_check=True)

# lif1 now has:
# input_type: {"input": array([128])}
# output_type: {"output": array([128])}

Conv2d with explicit entry point:

nodes = {
    "input": nir.Input(input_type={"input": np.array([3, 224, 224])}),
    "conv1": nir.Conv2d(
        weight=np.random.randn(64, 3, 7, 7).astype(np.float32),
        bias=np.zeros(64, dtype=np.float32),
        stride=(2, 2),
        padding=(3, 3)
    ),
    "lif1": nir.LIF(tau=0.01, r=1.0, v_leak=0.0, v_threshold=1.0),
    "output": nir.Output(output_type={"output": np.array([64, 112, 112])})
}

edges = [("input", "conv1"), ("conv1", "lif1"), ("lif1", "output")]
graph = nir.NIRGraph(nodes, edges, type_check=True)

# conv1.output_type computed as [64, 112, 112]
# lif1 infers shape from conv1

Recurrent connections:

nodes = {
    "input": nir.Input(input_type={"input": np.array([64])}),
    "lif": nir.LIF(tau=0.01, r=1.0, v_leak=0.0, v_threshold=1.0),
    "recurrent": nir.Affine(
        weight=np.random.randn(64, 64).astype(np.float32),
        bias=np.zeros(64, dtype=np.float32)
    ),
    "output": nir.Output(output_type={"output": np.array([64])})
}

# Cycle: lif -> recurrent -> lif
edges = [
    ("input", "lif"), 
    ("lif", "recurrent"), 
    ("recurrent", "lif"),  # Feedback
    ("lif", "output")
]

# SCC algorithm handles cycle correctly
graph = nir.NIRGraph(nodes, edges, type_check=True)

Test Suite

# tests/test_type_inference.py

import pytest
import numpy as np
import nir

class TestScalarParameterInference:
    
    def test_lif_scalar_params_with_predecessor(self):
        """LIF with scalar params infers shape from predecessor."""
        nodes = {
            "input": nir.Input(input_type={"input": np.array([128])}),
            "lif": nir.LIF(tau=0.01, r=1.0, v_leak=0.0, v_threshold=1.0),
            "output": nir.Output(output_type={"output": np.array([128])})
        }
        edges = [("input", "lif"), ("lif", "output")]
        
        graph = nir.NIRGraph(nodes, edges, type_check=True)
        
        assert np.array_equal(graph.nodes["lif"].input_type["input"], [128])
        assert np.array_equal(graph.nodes["lif"].output_type["output"], [128])
    
    def test_lif_scalar_params_no_predecessor_fails(self):
        """LIF with scalar params and no predecessor raises error."""
        nodes = {
            "lif": nir.LIF(tau=0.01, r=1.0, v_leak=0.0, v_threshold=1.0),
            "output": nir.Output(output_type={"output": np.array([128])})
        }
        edges = [("lif", "output")]
        
        with pytest.raises(nir.InferenceError):
            nir.NIRGraph(nodes, edges, type_check=True)
    
    def test_lif_explicit_params_override(self):
        """LIF with explicit array params uses those shapes."""
        nodes = {
            "input": nir.Input(input_type={"input": np.array([64])}),
            "lif": nir.LIF(
                tau=np.ones(128) * 0.01,  # Explicit 128 neurons
                r=np.ones(128),
                v_leak=np.zeros(128),
                v_threshold=np.ones(128)
            ),
            "output": nir.Output(output_type={"output": np.array([128])})
        }
        edges = [("input", "lif"), ("lif", "output")]
        
        # Should fail: input provides 64, LIF params say 128
        with pytest.raises(nir.InferenceError):
            nir.NIRGraph(nodes, edges, type_check=True)


class TestConv2dInference:
    
    def test_conv2d_computes_output_shape(self):
        """Conv2d computes output shape from input and kernel."""
        nodes = {
            "input": nir.Input(input_type={"input": np.array([3, 32, 32])}),
            "conv": nir.Conv2d(
                weight=np.random.randn(16, 3, 3, 3).astype(np.float32),
                bias=np.zeros(16, dtype=np.float32),
                stride=(1, 1),
                padding=(1, 1)
            ),
            "output": nir.Output(output_type={"output": np.array([16, 32, 32])})
        }
        edges = [("input", "conv"), ("conv", "output")]
        
        graph = nir.NIRGraph(nodes, edges, type_check=True)
        
        assert np.array_equal(graph.nodes["conv"].output_type["output"], [16, 32, 32])
    
    def test_conv2d_no_input_fails(self):
        """Conv2d without input_type or predecessor raises error."""
        nodes = {
            "conv": nir.Conv2d(
                weight=np.random.randn(16, 3, 3, 3).astype(np.float32),
                bias=np.zeros(16, dtype=np.float32)
            ),
            "output": nir.Output(output_type={"output": np.array([16, 30, 30])})
        }
        edges = [("conv", "output")]
        
        with pytest.raises(nir.InferenceError):
            nir.NIRGraph(nodes, edges, type_check=True)


class TestCyclicGraphs:
    
    def test_recurrent_connection(self):
        """Recurrent connection with typed entry point works."""
        nodes = {
            "input": nir.Input(input_type={"input": np.array([64])}),
            "lif": nir.LIF(tau=0.01, r=1.0, v_leak=0.0, v_threshold=1.0),
            "feedback": nir.Affine(
                weight=np.random.randn(64, 64).astype(np.float32),
                bias=np.zeros(64, dtype=np.float32)
            ),
            "output": nir.Output(output_type={"output": np.array([64])})
        }
        edges = [
            ("input", "lif"),
            ("lif", "feedback"),
            ("feedback", "lif"),
            ("lif", "output")
        ]
        
        graph = nir.NIRGraph(nodes, edges, type_check=True)
        
        assert graph.nodes["lif"].input_type is not None
        assert graph.nodes["feedback"].output_type is not None

[Jegp]

You asked whether it is worth the effort. For a 1024-neuron layer with 5 float32 parameters:

* Broadcast: 1024 x 5 x 4 = 20,480 bytes

* Scalar: 5 x 4 = 20 bytes

For a 10-layer network: 200KB vs 200 bytes :) The savings are significant for edge deployment, but I acknowledge this is a backend concern, not NIR's responsibility.

I see your argument. That's a lot of bytes. I also want to point out that HDF5 files actually do have compression available. I tried this on a few different files (I pushed the changes to #177 - your review would be appreciated :-) ). The results:

Network	Size compressed	Size uncompression
1x1 Linear (to compare baseline file size)	12K	8.1K
10000x10000 Linear with random weights	734M	763M
10000x10000 Linear with only ones	1.8M	763M

I'm surprised that a 1x1 Linear layer is 8Kb! But that's besides the point, I guess. Do you think this is a meaningful tradeoff? Meaning, if we can get that level of compression on large (broadcast-able) arrays, is it still worth the effort of doing the type inference?

Happy to hear your take!

[siddiquifaras]

Thanks for the quick turnaround on #177. I reviewed the compression changes and they look good. The gzip integration is clean.

To answer your question directly: yes, I still think type inference is worth it, but for a different reason than storage.

Compression and Type Inference Solve Different Problems

Problem	Compression	Type Inference
File size	Solves it	No effect
Scalar params failing type_check=True	No effect	Solves it
Runtime memory on edge devices	No effect	Enables backend optimization

The issue I originally filed was not primarily about file size. It was about this error:

ValueError: Type inference error: type mismatch: fc1.output: [[128]] -> lif.input: []

This happens because scalar parameters produce empty type annotations. Compression does not fix this. Users still cannot use scalar LIF parameters with type_check=True.

The Core Ask

What I am really asking for is:

# This should work without broadcasting
lif = nir.LIF(tau=0.01, r=1.0, v_leak=0.0, v_threshold=1.0)

# When placed in a graph with context
graph = nir.NIRGraph(
    nodes={"input": nir.Input(...), "fc1": nir.Affine(...), "lif": lif, ...},
    edges=[("input", "fc1"), ("fc1", "lif"), ...],
    type_check=True  # Should pass, not fail
)

Currently this fails. The only workaround is to broadcast scalars to arrays, which works but feels like fighting the API.

On Storage vs Runtime Memory

You are right that compression dramatically reduces file size for uniform arrays. That is valuable for distribution.

On the 200KB vs 200 bytes comparison I gave earlier: that was about runtime memory, not file size. When you load a compressed file, it decompresses to the full array in RAM. Compression helps storage and transfer. It does not help the edge device that has to hold those arrays in SRAM during inference.

This is admittedly a backend concern. But NIR could preserve the semantic information that "this is a shared scalar" rather than "these are 128 values that happen to be identical."

Proposal

Would you be open to a minimal PR that:

1. Adds resolve_input_type() / compute_output_type() methods to LIF only (as a proof of concept)
2. Hooks into the existing infer_types() flow
3. Includes tests

If it works well for LIF, we can extend to other nodes incrementally. If it causes issues, we revert.

Let me know if this would be welcome or if you would prefer to handle it differently.

Thanks for engaging on this.