Add option to drop unused inputs during build (#152)

Author: cbourjau

CHANGELOG.rst

@@ -7,6 +7,14 @@
 Change log
 ==========
 
+0.12.0 (2024-05-16)
+-------------------
+
+**New feature**
+
+- The :func:`spox.build` function gained the ``drop_unused_inputs`` argument.
+
+
 0.11.0 (2024-04-23)
 -------------------
 

pyproject.toml

@@ -72,4 +72,8 @@ check_untyped_defs = true
 
 [tool.pytest.ini_options]
 # This will be pytest's future default.
-addopts = "--import-mode=importlib"
+addopts = "--import-mode=importlib"
+filterwarnings = [
+    # Protobuf warning seen when running the test suite
+    "ignore:.*Type google.protobuf.pyext.*:DeprecationWarning:.*",
+]

src/spox/_public.py

@@ -54,7 +54,9 @@ def _temporary_renames(**kwargs: Var):
             arg._rename(name)
 
 
-def build(inputs: Dict[str, Var], outputs: Dict[str, Var]) -> onnx.ModelProto:
+def build(
+    inputs: Dict[str, Var], outputs: Dict[str, Var], *, drop_unused_inputs=False
+) -> onnx.ModelProto:
     """
     Builds an ONNX Model with given model inputs and outputs.
 
@@ -71,6 +73,11 @@ def build(inputs: Dict[str, Var], outputs: Dict[str, Var]) -> onnx.ModelProto:
         Model outputs. Keys are names, values may be any ``Var``.
         Building will resolve what nodes were used in the construction
         of output variables.
+    drop_unused_inputs
+        If ``False``, only inputs that are needed for the computation
+        of the ``outputs`` will appear as inputs of the returned
+        ``ModelProto``. Otherwise, all inputs are required by the
+        returned object (default).
 
     Returns
     -------
@@ -81,6 +88,11 @@ def build(inputs: Dict[str, Var], outputs: Dict[str, Var]) -> onnx.ModelProto:
         the newest one. The minimum ``ai.onnx`` version is set to 14 to
         avoid tooling issues with legacy versions.
 
+    Raises
+    ------
+    KeyError
+        If the ``outputs`` cannot be build from the given ``inputs``.
+
     Examples
     --------
     >>> import numpy as np
@@ -113,8 +125,15 @@ def build(inputs: Dict[str, Var], outputs: Dict[str, Var]) -> onnx.ModelProto:
 
     with _temporary_renames(**inputs):
         graph = results(**outputs)
-        graph = graph.with_arguments(*inputs.values())
-        return graph.to_onnx_model()
+        if not drop_unused_inputs:
+            graph = graph.with_arguments(*inputs.values())
+        model_proto = graph.to_onnx_model()
+
+    # Validate that no further inputs were required.
+    if any(inp.name not in inputs for inp in model_proto.graph.input):
+        raise KeyError("Model requires additional inputs not provided in 'inputs'.")
+
+    return model_proto
 
 
 class _InlineCall(Protocol):

tests/test_public.py

@@ -100,3 +100,22 @@ def test_shallow_deepcopy_var_raises():
     a = argument(Tensor(float, ()))
     with pytest.raises(ValueError):
         deepcopy(a)
+
+
+def test_build_drop_unused_arguments():
+    a = argument(Tensor(float, ()))
+    unused = argument(Tensor(float, ()))
+    c = op.add(a, a)
+    model_proto = build({"a": a, "unused": unused}, {"c": c}, drop_unused_inputs=True)
+    actual_inputs = [el.name for el in model_proto.graph.input]
+
+    assert actual_inputs == ["a"]
+
+
+@pytest.mark.parametrize("drop_unused", [True, False])
+def test_raise_missing_input(drop_unused):
+    a = argument(Tensor(float, ()))
+    b = argument(Tensor(float, ()))
+
+    with pytest.raises(KeyError):
+        build({"a": a}, {"c": op.add(a, b)}, drop_unused_inputs=drop_unused)