Tests for constant.

Author: MKaptein

sclblonnx/constant.py

@@ -2,7 +2,7 @@
 from onnx import helper as xhelp
 from onnx import onnx_ml_pb2 as xpb2
 
-from sclblonnx.utils import _data_type
+from sclblonnx.utils import _data_type, _print
 from sclblonnx.node import add_node
 
 
@@ -17,10 +17,15 @@ def constant(name: str,
         name: Name of the (output value of the) constant node to determine the graph topology
         value: Values of the node (as a np.array)
         data_type: Data type of the node
+        **kwargs
 
     Returns:
-        The extended graph.
+        A constant node.
     """
+    if not name:
+        _print("Unable to create unnamed constant.")
+        return False
+
     dtype = _data_type(data_type)
     if not dtype:
         return False
@@ -30,7 +35,7 @@ def constant(name: str,
                                         value=xhelp.make_tensor(name=name + "-values", data_type=dtype,
                                                                 dims=value.shape, vals=value.flatten()), **kwargs)
     except Exception as e:
-        print("Unable to create the constant node: " + str(e))
+        _print("Unable to create the constant node: " + str(e))
         return False
 
     return constant_node
@@ -43,7 +48,9 @@ def add_constant(
         value: np.array,
         data_type: str,
         **kwargs):
-    """ Add a constant node to an existing graph
+    """ Create and add a constant node to an existing graph.
+
+    Note: use add_node() if you want to add an existing constant node to an existing graph
 
     Args:
         graph: A graph, onnx.onnx_ml_pb2.GraphProto.
@@ -67,12 +74,17 @@ def add_constant(
                                         value=xhelp.make_tensor(name=name + "-values", data_type=dtype,
                                                                 dims=value.shape, vals=value.flatten()), **kwargs)
     except Exception as e:
-        print("Unable to create the constant node: " + str(e))
+        _print("Unable to create the constant node: " + str(e))
+        return False
+
+    try:
+        graph = add_node(graph, constant_node, **kwargs)
+    except Exception as e:
+        _print("Unable to add the constant node to the graph: " + str(e))
         return False
 
-    graph = add_node(graph, constant_node, **kwargs)
     if not graph:
-        print("Unable to add constant node to graph.")
+        _print("Unable to add constant node to graph.")
         return False
     return graph
 

sclblonnx/utils.py

@@ -145,7 +145,7 @@ def _data_type(data_string: str):
     for key, val in glob.DATA_TYPES.items():
         if key == data_string:
             return val
-    print("Data type not found. Use `list_data_types()` to list all supported data types.")
+    _print("Data type not found. Use `list_data_types()` to list all supported data types.")
     return False
 
 

test/.tmp.onnx

@@ -0,0 +1,32 @@
+from sclblonnx import constant, empty_graph, node, add_node, add_constant, add_output, run, display, clean
+import numpy as np
+
+
+def test_constant():
+    c = constant("", np.array([1,2,]), "FLOAT")
+    assert not c, "Constant creation should have failed without a name."
+    c = constant("constant", np.array([1,2,]), "NONE")
+    assert not c, "Constant creation should have failed without a valid data type."
+    c = constant("constant", np.array([1,2,]), "FLOAT")
+    check = getattr(c, "output", False)
+    assert check[0] == "constant", "Constant creation should have worked."
+
+
+def test_add_constant():
+
+    # Simple add graph
+    g = empty_graph()
+    n1 = node('Add', inputs=['x1', 'x2'], outputs=['sum'])
+    g = add_node(g, n1)
+
+    # Add input and constant
+    g = add_constant(g, 'x1', np.array([1]), "INT64")
+    g = add_constant(g, 'x2', np.array([5]), "INT64")
+
+    # Output:
+    g = add_output(g, 'sum', "INT64", [1])
+
+    # This works, but seems to fail for other data types...
+    result = run(g, inputs={}, outputs=["sum"])
+    assert result[0] == 6, "Add constant failed."
+    # todo(McK): Does not work for INT16 / INT8, check?