Skip to content

fix: resolve scalar parameter types from graph context during inference #189

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
siddiquifaras wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix: resolve scalar parameter types from graph context during inference #189

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a2c7fa7
fix: resolve scalar parameter types from graph context during inference
siddiquifaras Feb 25, 2026
0dce00d
style: apply black formatting to graph.py
siddiquifaras Mar 9, 2026
7349584
test: verify port naming pattern for scalar type inference
siddiquifaras Mar 9, 2026
c6b7284
style: apply black formatting to test files
siddiquifaras Mar 9, 2026
7d9360d
fix: resolve merge conflict with main (structural validation tests)
siddiquifaras Mar 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 29 additions & 9 deletions nir/ir/graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -360,16 +360,36 @@ def infer_types(self):
for k, v in pre_node.output_type.items()
}
elif type_mismatch:
# set post input_type to be the same as pre output_type
pre_repr = (
f"{pre_key}.output: {np.array(list(pre_node.output_type.values()))}"
)
post_repr = (
f"{post_key}.input: {np.array(list(post_node.input_type.values()))}"
)
raise ValueError(
f"Type inference error: type mismatch: {pre_repr} -> {post_repr}"
# Check if post node has scalar-derived types (empty shape arrays).
# Scalar parameters (0-d arrays) produce input_type/output_type with
# empty arrays via np.array(param.shape) where shape=(). Resolve these
# by adopting the predecessor's output type. Parameters are untouched.
post_values = list(post_node.input_type.values())
is_scalar_type = all(
isinstance(v, np.ndarray) and v.size == 0 for v in post_values
)
if is_scalar_type:
post_node.input_type = {
k.replace("output", "input"): v.copy()
Copy link
Copy Markdown
Collaborator

@Jegp Jegp Mar 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me make sure I'm understanding this: you're taking the dict pre_node.output.items(), and then replacing the "output" node with "input" copying v as the value. Is this a good strategy? I get the idea that we want to copy the output parameter, but, in principle, one node can have multiple outputs.

I'm not sure I have a great solution for this, but it's worth thinking a bit about. Would it be an option to build a new dictionary that only aligns the input? It's a bit of a headache because the ports are slightly underspecified (i. e. people can (ab)use them in whichever way), but I'd like to avoid constraining that too much if it's not necessary.

Copy link
Copy Markdown
Collaborator Author

@siddiquifaras siddiquifaras Mar 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the feedback. The k.replace("output", "input") pattern is consistent with the existing implementation in infer_types() (lines 316–318), which uses the same approach for undefined input types.

1. Current NIR implementation

All nodes use single-port naming ({"input": ...}, {"output": ...}), and multi port nodes currently raise NotImplementedError in check_types() (lines 216–218). In this context, the string replacement works reliably.

2. Future multi port support

When multi port nodes are introduced, NIR edges connect entire nodes (all ports to all ports), not individual ports. The type check at line 307 enforces port count consistency:

len(post_node.input_type) != len(pre_node.output_type)

Copying the full output_type dict to input_type therefore remains correct.

3. Port naming convention

The approach also extends to indexed ports (e.g., output_0 → input_0, output_1 → input_1) as long as the naming convention retains the "output" / "input" substrings.

i've added a test (test_scalar_lif_port_naming) to confirm that port naming is preserved and to document this behavior for future multi port support.

for k, v in pre_node.output_type.items()
}
post_node.output_type = {
k.replace("input", "output"): v.copy()
for k, v in post_node.input_type.items()
}
else:
pre_repr = (
f"{pre_key}.output: "
f"{np.array(list(pre_node.output_type.values()))}"
)
post_repr = (
f"{post_key}.input: "
f"{np.array(list(post_node.input_type.values()))}"
)
raise ValueError(
f"Type inference error: type mismatch: "
f"{pre_repr} -> {post_repr}"
)

# make sure that output nodes have output_type = input_type
if isinstance(post_node, Output):
Expand Down
17 changes: 10 additions & 7 deletions nir/serialization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,13 +242,16 @@ def write_recursive(group: h5py.Group, node: dict) -> None:
elif isinstance(v, str):
group.create_dataset(k, data=v, dtype=h5py.string_dtype())
elif isinstance(v, np.ndarray):
group.create_dataset(
k,
data=v,
dtype=v.dtype,
compression=compression,
compression_opts=compression_opts,
)
if v.ndim > 0:
group.create_dataset(
k,
data=v,
dtype=v.dtype,
compression=compression,
compression_opts=compression_opts,
)
else:
group.create_dataset(k, data=v, dtype=v.dtype)
elif isinstance(v, dict):
write_recursive(group.create_group(str(k)), v)
else:
Expand Down
233 changes: 233 additions & 0 deletions tests/test_ir.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -755,6 +755,239 @@ def test_type_check_recurrent():
)


def test_scalar_lif_type_inference():
"""Scalar LIF parameters get types resolved from Input predecessor."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([64])),
"lif": nir.LIF(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
),
"output": nir.Output(np.array([64])),
},
edges=[("input", "lif"), ("lif", "output")],
)
assert np.array_equal(graph.nodes["lif"].input_type["input"], [64])
assert np.array_equal(graph.nodes["lif"].output_type["output"], [64])
assert graph.nodes["lif"].tau.shape == ()
assert graph.nodes["lif"].r.shape == ()


def test_scalar_lif_with_affine_predecessor():
"""Scalar LIF infers shape from Affine predecessor output."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([784])),
"fc": nir.Affine(
weight=np.random.randn(128, 784),
bias=np.zeros(128),
),
"lif": nir.LIF(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
),
"output": nir.Output(np.array([128])),
},
edges=[("input", "fc"), ("fc", "lif"), ("lif", "output")],
)
assert np.array_equal(graph.nodes["lif"].input_type["input"], [128])
assert np.array_equal(graph.nodes["lif"].output_type["output"], [128])


def test_scalar_cubalif_type_inference():
"""Scalar CubaLIF parameters get types resolved from predecessor."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([32])),
"lif": nir.CubaLIF(
tau_syn=np.array(0.01),
tau_mem=np.array(0.02),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
),
"output": nir.Output(np.array([32])),
},
edges=[("input", "lif"), ("lif", "output")],
)
assert np.array_equal(graph.nodes["lif"].input_type["input"], [32])
assert np.array_equal(graph.nodes["lif"].output_type["output"], [32])


def test_scalar_if_type_inference():
"""Scalar IF parameters get types resolved from predecessor."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([16])),
"neuron": nir.IF(
r=np.array(1.0),
v_threshold=np.array(1.0),
),
"output": nir.Output(np.array([16])),
},
edges=[("input", "neuron"), ("neuron", "output")],
)
assert np.array_equal(graph.nodes["neuron"].input_type["input"], [16])
assert np.array_equal(graph.nodes["neuron"].output_type["output"], [16])


def test_scalar_li_type_inference():
"""Scalar LI parameters get types resolved from predecessor."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([8])),
"neuron": nir.LI(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
),
"output": nir.Output(np.array([8])),
},
edges=[("input", "neuron"), ("neuron", "output")],
)
assert np.array_equal(graph.nodes["neuron"].input_type["input"], [8])
assert np.array_equal(graph.nodes["neuron"].output_type["output"], [8])


def test_scalar_lif_multi_layer():
"""Multiple scalar LIF layers get types resolved sequentially."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([784])),
"fc1": nir.Affine(weight=np.random.randn(128, 784), bias=np.zeros(128)),
"lif1": nir.LIF(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
),
"fc2": nir.Affine(weight=np.random.randn(10, 128), bias=np.zeros(10)),
"lif2": nir.LIF(
tau=np.array(0.02),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(0.5),
),
"output": nir.Output(np.array([10])),
},
edges=[
("input", "fc1"),
("fc1", "lif1"),
("lif1", "fc2"),
("fc2", "lif2"),
("lif2", "output"),
],
)
assert np.array_equal(graph.nodes["lif1"].input_type["input"], [128])
assert np.array_equal(graph.nodes["lif1"].output_type["output"], [128])
assert np.array_equal(graph.nodes["lif2"].input_type["input"], [10])
assert np.array_equal(graph.nodes["lif2"].output_type["output"], [10])


def test_scalar_lif_recurrent():
"""Scalar LIF works in a recurrent graph with feedback connection."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([64])),
"lif": nir.LIF(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
),
"feedback": nir.Linear(weight=np.random.randn(64, 64)),
"output": nir.Output(np.array([64])),
},
edges=[
("input", "lif"),
("lif", "feedback"),
("feedback", "lif"),
("lif", "output"),
],
)
assert np.array_equal(graph.nodes["lif"].input_type["input"], [64])
assert np.array_equal(graph.nodes["lif"].output_type["output"], [64])


def test_nonscalar_mismatch_still_fails():
"""Non-scalar type mismatches still raise ValueError as before."""
with pytest.raises(ValueError, match="type mismatch"):
nir.NIRGraph(
nodes={
"input": nir.Input(np.array([64])),
"lif": nir.LIF(
tau=np.ones(128) * 0.01,
r=np.ones(128),
v_leak=np.zeros(128),
v_threshold=np.ones(128),
),
"output": nir.Output(np.array([128])),
},
edges=[("input", "lif"), ("lif", "output")],
)


def test_explicit_array_params_unchanged():
"""Nodes with explicit (non-scalar) array params behave as before."""
graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([3])),
"lif": nir.LIF(
tau=np.array([0.01, 0.02, 0.03]),
r=np.array([1.0, 1.0, 1.0]),
v_leak=np.array([0.0, 0.0, 0.0]),
v_threshold=np.array([1.0, 1.0, 1.0]),
),
"output": nir.Output(np.array([3])),
},
edges=[("input", "lif"), ("lif", "output")],
)
assert np.array_equal(graph.nodes["lif"].input_type["input"], [3])
assert np.array_equal(graph.nodes["lif"].output_type["output"], [3])


def test_scalar_lif_port_naming():
"""Verify that scalar type inference preserves port naming conventions.

This test demonstrates that the string replacement pattern (output -> input)
works correctly for both single-port and hypothetical multi-port nodes.
When multi-port support is added to NIR, port names like 'output_0', 'output_1'
will correctly map to 'input_0', 'input_1'.
"""
affine = nir.Affine(weight=np.random.randn(10, 5), bias=np.zeros(10))
scalar_lif = nir.LIF(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
)

assert "input" in scalar_lif.input_type
assert scalar_lif.input_type["input"].size == 0
assert "output" in scalar_lif.output_type
assert scalar_lif.output_type["output"].size == 0

graph = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([5])),
"affine": affine,
"lif": scalar_lif,
"output": nir.Output(np.array([10])),
},
edges=[("input", "affine"), ("affine", "lif"), ("lif", "output")],
)

assert "input" in graph.nodes["lif"].input_type
assert np.array_equal(graph.nodes["lif"].input_type["input"], [10])
assert "output" in graph.nodes["lif"].output_type
assert np.array_equal(graph.nodes["lif"].output_type["output"], [10])


def test_validate_structure_dangling_source():
"""Edge referencing a non-existent source node raises ValueError."""
with pytest.raises(ValueError, match="does not exist"):
Expand Down
65 changes: 65 additions & 0 deletions tests/test_readwrite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,6 +331,71 @@ def test_read_without_type_check():
)


def test_scalar_lif_readwrite():
"""Scalar LIF graph survives write/read roundtrip with type_check=True."""
ir = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([64])),
"lif": nir.LIF(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
),
"output": nir.Output(np.array([64])),
},
edges=[("input", "lif"), ("lif", "output")],
)
assert np.array_equal(ir.nodes["lif"].input_type["input"], [64])

with tempfile.TemporaryFile() as fp:
nir.write(fp, ir)
ir2 = nir.read(fp)

assert np.array_equal(ir2.nodes["lif"].input_type["input"], [64])
assert np.array_equal(ir2.nodes["lif"].output_type["output"], [64])
assert ir2.nodes["lif"].tau.shape == ()
assert ir2.nodes["lif"].r.shape == ()


def test_scalar_lif_multilayer_readwrite():
"""Multi-layer scalar LIF graph survives write/read roundtrip."""
ir = nir.NIRGraph(
nodes={
"input": nir.Input(np.array([784])),
"fc1": nir.Affine(weight=np.random.randn(128, 784), bias=np.zeros(128)),
"lif1": nir.LIF(
tau=np.array(0.01),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(1.0),
),
"fc2": nir.Affine(weight=np.random.randn(10, 128), bias=np.zeros(10)),
"lif2": nir.LIF(
tau=np.array(0.02),
r=np.array(1.0),
v_leak=np.array(0.0),
v_threshold=np.array(0.5),
),
"output": nir.Output(np.array([10])),
},
edges=[
("input", "fc1"),
("fc1", "lif1"),
("lif1", "fc2"),
("fc2", "lif2"),
("lif2", "output"),
],
)

with tempfile.TemporaryFile() as fp:
nir.write(fp, ir)
ir2 = nir.read(fp)

assert np.array_equal(ir2.nodes["lif1"].input_type["input"], [128])
assert np.array_equal(ir2.nodes["lif2"].input_type["input"], [10])


def test_serialize_deserialize_data():
graph_data = nir.NIRGraphData(
nodes={
Expand Down
Loading