Don't include unused inputs in ndonnx_schema when they are dropped by ndonnx.build (#211)

Author: sgilm

CHANGELOG.rst

@@ -7,6 +7,14 @@
 Changelog
 =========
 
+0.18.2 (unreleased)
+-------------------
+
+**Bug fix**
+
+- Fixed an issue whereby any unused inputs dropped via the `drop_unused` argument to :func:`ndonnx.build` were still included in `ndonnx_schema`.
+
+
 0.18.1 (2026-04-22)
 -------------------
 

ndonnx/_build.py

@@ -1,7 +1,6 @@
-# Copyright (c) QuantCo 2023-2025
+# Copyright (c) QuantCo 2023-2026
 # SPDX-License-Identifier: BSD-3-Clause
 
-
 import onnx
 from spox import Var
 from spox import build as spox_build
@@ -32,9 +31,19 @@ def build(
 
     mp = spox_build(ins, outs, drop_unused_inputs=drop_unused)
 
+    # Find the names of inputs that have not been dropped, so that only these
+    # are inserted into the schema.
+    graph_inputs = [input.name for input in mp.graph.input]
+    kept_inputs = []
+    for name, arr in inputs.items():
+        if any(n in graph_inputs for n in _disassemble_named_array(name, arr)):
+            kept_inputs.append(name)
+
     schema_v1 = {
         "ndonnx_schema": SchemaV1(
-            input_schema={k: v.dtype.__ndx_infov1__ for k, v in inputs.items()},
+            input_schema={
+                k: v.dtype.__ndx_infov1__ for k, v in inputs.items() if k in kept_inputs
+            },
             output_schema={k: v.dtype.__ndx_infov1__ for k, v in outputs.items()},
             version=1,
         ).to_json()
@@ -45,14 +54,21 @@ def build(
 
 
 def _arrays_to_vars(dct_of_arrs: dict[str, Array]) -> dict[str, Var]:
+    out: dict[str, Var] = {}
+    for name, arr in dct_of_arrs.items():
+        out |= _disassemble_named_array(name, arr)
+    return out
+
+
+def _disassemble_named_array(name: str, arr: Array) -> dict[str, Var]:
+    # Take an Array, and create a map of its component parts prefixed with a name.
     # TODO: Use a different separator for the public name and the nested components?
     public_separator = "_"
-    out = {}
-    for k, v in dct_of_arrs.items():
-        components = v._tyarray.disassemble()
-        if isinstance(components, Var):
-            out[k] = components
-            continue
-        for k_inner, v_inner in components.items():
-            out[f"{k}{public_separator}{k_inner}"] = v_inner
+    out: dict[str, Var] = {}
+    components = arr._tyarray.disassemble()
+    if isinstance(components, Var):
+        out[name] = components
+    else:
+        for k, v in components.items():
+            out[f"{name}{public_separator}{k}"] = v
     return out

tests/test_build_utils.py

@@ -1,9 +1,10 @@
-# Copyright (c) QuantCo 2023-2025
+# Copyright (c) QuantCo 2023-2026
 # SPDX-License-Identifier: BSD-3-Clause
 import json
 from pathlib import Path
 
 import pytest
+from onnx import ModelProto
 
 import ndonnx as ndx
 import ndonnx._typed_array as tydx
@@ -106,3 +107,39 @@ def test_schema_against_snapshots(dtype, update_schema_snapshots):
     # test json round trip of schema data
     assert candidate_schemas["input_schema"]["a"] == a.dtype.__ndx_infov1__.__dict__
     assert candidate_schemas["output_schema"]["b"] == b.dtype.__ndx_infov1__.__dict__
+
+
+def _get_input_schema(model_proto: ModelProto):
+    for entry in model_proto.metadata_props:
+        if entry.key == "ndonnx_schema":
+            return json.loads(entry.value)["input_schema"]
+    raise KeyError("'ndonnx_schema' not present")
+
+
+def test_unused_inputs_not_in_schema():
+    # Case 1: simple graph with an unused input.
+    a = ndx.argument(shape=("N",), dtype=ndx.float64)
+    b = ndx.argument(shape=("N",), dtype=ndx.float64)
+    unused = ndx.argument(shape=("N",), dtype=ndx.float64)
+
+    c = a + b
+
+    no_drop = ndx.build(
+        inputs={"a": a, "b": b, "unused": unused}, outputs={"c": c}, drop_unused=False
+    )
+
+    assert "unused" in _get_input_schema(no_drop)
+
+    drop = ndx.build(
+        inputs={"a": a, "b": b, "unused": unused}, outputs={"c": c}, drop_unused=True
+    )
+    assert "unused" not in _get_input_schema(drop)
+
+    # Case 2: part of a nullable input is not used.
+    a = ndx.argument(shape=("N",), dtype=ndx.nint64)
+
+    b = ndx.isnan(a)
+
+    drop = ndx.build(inputs={"a": a}, outputs={"b": b}, drop_unused=True)
+
+    assert "a" in _get_input_schema(no_drop)