# Copyright (c) ONNX Project Contributors # SPDX-License-Identifier: Apache-2.0 import itertools import random import struct import unittest from typing import Any, List, Tuple import numpy as np import parameterized import pytest import version_utils from onnx import ( AttributeProto, GraphProto, ModelProto, OptionalProto, SequenceProto, TensorProto, TypeProto, checker, defs, helper, numpy_helper, ) from onnx.reference.op_run import to_array_extended class TestHelperAttributeFunctions(unittest.TestCase): def test_attr_float(self) -> None: # float attr = helper.make_attribute("float", 1.0) self.assertEqual(attr.name, "float") self.assertEqual(attr.f, 1.0) checker.check_attribute(attr) # float with scientific attr = helper.make_attribute("float", 1e10) self.assertEqual(attr.name, "float") self.assertEqual(attr.f, 1e10) checker.check_attribute(attr) def test_attr_int(self) -> None: # integer attr = helper.make_attribute("int", 3) self.assertEqual(attr.name, "int") self.assertEqual(attr.i, 3) checker.check_attribute(attr) # long integer attr = helper.make_attribute("int", 5) self.assertEqual(attr.name, "int") self.assertEqual(attr.i, 5) checker.check_attribute(attr) # octinteger attr = helper.make_attribute("int", 0o1701) self.assertEqual(attr.name, "int") self.assertEqual(attr.i, 0o1701) checker.check_attribute(attr) # hexinteger attr = helper.make_attribute("int", 0x1701) self.assertEqual(attr.name, "int") self.assertEqual(attr.i, 0x1701) checker.check_attribute(attr) def test_attr_doc_string(self) -> None: attr = helper.make_attribute("a", "value") self.assertEqual(attr.name, "a") self.assertEqual(attr.doc_string, "") attr = helper.make_attribute("a", "value", "doc") self.assertEqual(attr.name, "a") self.assertEqual(attr.doc_string, "doc") def test_attr_string(self) -> None: # bytes attr = helper.make_attribute("str", b"test") self.assertEqual(attr.name, "str") self.assertEqual(attr.s, b"test") checker.check_attribute(attr) # unspecified attr = helper.make_attribute("str", "test") self.assertEqual(attr.name, "str") self.assertEqual(attr.s, b"test") checker.check_attribute(attr) # unicode attr = helper.make_attribute("str", "test") self.assertEqual(attr.name, "str") self.assertEqual(attr.s, b"test") checker.check_attribute(attr) # empty str attr = helper.make_attribute("str", "") self.assertEqual(attr.name, "str") self.assertEqual(helper.get_attribute_value(attr), b"") checker.check_attribute(attr) def test_attr_repeated_float(self) -> None: attr = helper.make_attribute("floats", [1.0, 2.0]) self.assertEqual(attr.name, "floats") self.assertEqual(list(attr.floats), [1.0, 2.0]) checker.check_attribute(attr) def test_attr_repeated_int(self) -> None: attr = helper.make_attribute("ints", [1, 2]) self.assertEqual(attr.name, "ints") self.assertEqual(list(attr.ints), [1, 2]) checker.check_attribute(attr) def test_attr_repeated_mixed_floats_and_ints(self) -> None: attr = helper.make_attribute("mixed", [1, 2, 3.0, 4.5]) self.assertEqual(attr.name, "mixed") self.assertEqual(list(attr.floats), [1.0, 2.0, 3.0, 4.5]) checker.check_attribute(attr) def test_attr_repeated_str(self) -> None: attr = helper.make_attribute("strings", ["str1", "str2"]) self.assertEqual(attr.name, "strings") self.assertEqual(list(attr.strings), [b"str1", b"str2"]) checker.check_attribute(attr) def test_attr_repeated_tensor_proto(self) -> None: tensors = [ helper.make_tensor( name="a", data_type=TensorProto.FLOAT, dims=(1,), vals=np.ones(1) ), helper.make_tensor( name="b", data_type=TensorProto.FLOAT, dims=(1,), vals=np.ones(1) ), ] attr = helper.make_attribute("tensors", tensors) self.assertEqual(attr.name, "tensors") self.assertEqual(list(attr.tensors), tensors) checker.check_attribute(attr) def test_attr_sparse_tensor_proto(self) -> None: dense_shape = [3, 3] sparse_values = [1.764052391052246, 0.40015721321105957, 0.978738009929657] values_tensor = helper.make_tensor( name="sparse_values", data_type=TensorProto.FLOAT, dims=[len(sparse_values)], vals=np.array(sparse_values).astype(np.float32), raw=False, ) linear_indices = [2, 3, 5] indices_tensor = helper.make_tensor( name="indices", data_type=TensorProto.INT64, dims=[len(linear_indices)], vals=np.array(linear_indices).astype(np.int64), raw=False, ) sparse_tensor = helper.make_sparse_tensor( values_tensor, indices_tensor, dense_shape ) attr = helper.make_attribute("sparse_attr", sparse_tensor) self.assertEqual(attr.name, "sparse_attr") checker.check_sparse_tensor(helper.get_attribute_value(attr)) checker.check_attribute(attr) def test_attr_sparse_tensor_repeated_protos(self) -> None: dense_shape = [3, 3] sparse_values = [1.764052391052246, 0.40015721321105957, 0.978738009929657] values_tensor = helper.make_tensor( name="sparse_values", data_type=TensorProto.FLOAT, dims=[len(sparse_values)], vals=np.array(sparse_values).astype(np.float32), raw=False, ) linear_indices = [2, 3, 5] indices_tensor = helper.make_tensor( name="indices", data_type=TensorProto.INT64, dims=[len(linear_indices)], vals=np.array(linear_indices).astype(np.int64), raw=False, ) sparse_tensor = helper.make_sparse_tensor( values_tensor, indices_tensor, dense_shape ) repeated_sparse = [sparse_tensor, sparse_tensor] attr = helper.make_attribute("sparse_attrs", repeated_sparse) self.assertEqual(attr.name, "sparse_attrs") checker.check_attribute(attr) for s in helper.get_attribute_value(attr): checker.check_sparse_tensor(s) def test_attr_repeated_graph_proto(self) -> None: graphs = [GraphProto(), GraphProto()] graphs[0].name = "a" graphs[1].name = "b" attr = helper.make_attribute("graphs", graphs) self.assertEqual(attr.name, "graphs") self.assertEqual(list(attr.graphs), graphs) checker.check_attribute(attr) def test_attr_type_proto(self) -> None: # type_proto type_proto = TypeProto() attr = helper.make_attribute("type_proto", type_proto) self.assertEqual(attr.name, "type_proto") self.assertEqual(attr.tp, type_proto) self.assertEqual(attr.type, AttributeProto.TYPE_PROTO) # type_protos types = [TypeProto(), TypeProto()] attr = helper.make_attribute("type_protos", types) self.assertEqual(attr.name, "type_protos") self.assertEqual(list(attr.type_protos), types) self.assertEqual(attr.type, AttributeProto.TYPE_PROTOS) def test_attr_empty_list(self) -> None: attr = helper.make_attribute("empty", [], attr_type=AttributeProto.STRINGS) self.assertEqual(attr.type, AttributeProto.STRINGS) self.assertEqual(len(attr.strings), 0) self.assertRaises(ValueError, helper.make_attribute, "empty", []) def test_attr_mismatch(self) -> None: with self.assertRaisesRegex(TypeError, "Inferred attribute type 'FLOAT'"): helper.make_attribute("test", 6.4, attr_type=AttributeProto.STRING) def test_is_attr_legal(self) -> None: # no name, no field attr = AttributeProto() self.assertRaises(checker.ValidationError, checker.check_attribute, attr) # name, but no field attr = AttributeProto() attr.name = "test" self.assertRaises(checker.ValidationError, checker.check_attribute, attr) # name, with two fields attr = AttributeProto() attr.name = "test" attr.f = 1.0 attr.i = 2 self.assertRaises(checker.ValidationError, checker.check_attribute, attr) def test_is_attr_legal_verbose(self) -> None: def _set( attr: AttributeProto, type_: AttributeProto.AttributeType, var: str, value: Any, ) -> None: setattr(attr, var, value) attr.type = type_ def _extend( attr: AttributeProto, type_: AttributeProto.AttributeType, var: List[Any], value: Any, ) -> None: var.extend(value) attr.type = type_ SET_ATTR = [ (lambda attr: _set(attr, AttributeProto.FLOAT, "f", 1.0)), (lambda attr: _set(attr, AttributeProto.INT, "i", 1)), (lambda attr: _set(attr, AttributeProto.STRING, "s", b"str")), ( lambda attr: _extend( attr, AttributeProto.FLOATS, attr.floats, [1.0, 2.0] ) ), (lambda attr: _extend(attr, AttributeProto.INTS, attr.ints, [1, 2])), ( lambda attr: _extend( attr, AttributeProto.STRINGS, attr.strings, [b"a", b"b"] ) ), ] # Randomly set one field, and the result should be legal. for _i in range(100): attr = AttributeProto() attr.name = "test" random.choice(SET_ATTR)(attr) checker.check_attribute(attr) # Randomly set two fields, and then ensure helper function catches it. for _i in range(100): attr = AttributeProto() attr.name = "test" for func in random.sample(SET_ATTR, 2): func(attr) self.assertRaises(checker.ValidationError, checker.check_attribute, attr) class TestHelperNodeFunctions(unittest.TestCase): def test_node_no_arg(self) -> None: self.assertTrue(defs.has("Relu")) node_def = helper.make_node("Relu", ["X"], ["Y"], name="test") self.assertEqual(node_def.op_type, "Relu") self.assertEqual(node_def.name, "test") self.assertEqual(list(node_def.input), ["X"]) self.assertEqual(list(node_def.output), ["Y"]) def test_attr_doc_string(self) -> None: node_def = helper.make_node("Relu", ["X"], ["Y"], name="test", doc_string="doc") self.assertEqual(node_def.doc_string, "doc") def test_node_with_arg(self) -> None: self.assertTrue(defs.has("Relu")) # Note: Relu actually does not need an arg, but let's # test it. node_def = helper.make_node("Relu", ["X"], ["Y"], arg_value=1) self.assertEqual(node_def.op_type, "Relu") self.assertEqual(list(node_def.input), ["X"]) self.assertEqual(list(node_def.output), ["Y"]) self.assertEqual(len(node_def.attribute), 1) self.assertEqual(node_def.attribute[0], helper.make_attribute("arg_value", 1)) def test_node_domain(self) -> None: node_def = helper.make_node( "Relu", ["X"], ["Y"], name="test", doc_string="doc", domain="test.domain" ) self.assertEqual(node_def.domain, "test.domain") def test_graph(self) -> None: node_def1 = helper.make_node("Relu", ["X"], ["Y"]) node_def2 = helper.make_node("Add", ["X", "Y"], ["Z"]) value_info = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2])] graph = helper.make_graph( [node_def1, node_def2], "test", [helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2])], [helper.make_tensor_value_info("Z", TensorProto.FLOAT, [1, 2])], doc_string=None, value_info=value_info, ) self.assertEqual(graph.name, "test") self.assertEqual(len(graph.node), 2) self.assertEqual(graph.node[0], node_def1) self.assertEqual(graph.node[1], node_def2) self.assertEqual(graph.doc_string, "") self.assertEqual(graph.value_info[0], value_info[0]) def test_graph_docstring(self) -> None: graph = helper.make_graph([], "my graph", [], [], None, "my docs") self.assertEqual(graph.name, "my graph") self.assertEqual(graph.doc_string, "my docs") def test_model(self) -> None: node_def = helper.make_node("Relu", ["X"], ["Y"]) graph_def = helper.make_graph( [node_def], "test", [helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2])], [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2])], ) self.assertRaises(AttributeError, helper.make_model, graph_def, xxx=1) model_def = helper.make_model(graph_def, producer_name="test") self.assertEqual(model_def.producer_name, "test") def test_model_docstring(self) -> None: graph = helper.make_graph([], "my graph", [], []) model_def = helper.make_model(graph, doc_string="test") # models may have their own documentation, but don't have a name # their name is the domain-qualified name of the underlying graph. self.assertFalse(hasattr(model_def, "name")) self.assertEqual(model_def.doc_string, "test") def test_model_metadata_props(self) -> None: graph = helper.make_graph([], "my graph", [], []) model_def = helper.make_model(graph, doc_string="test") helper.set_model_props( model_def, {"Title": "my graph", "Keywords": "test;graph"} ) checker.check_model(model_def) helper.set_model_props( model_def, {"Title": "my graph", "Keywords": "test;graph"} ) checker.check_model(model_def) # helper replaces, so no dupe dupe = model_def.metadata_props.add() dupe.key = "Title" dupe.value = "Other" self.assertRaises(checker.ValidationError, checker.check_model, model_def) def test_model_irversion(self) -> None: def mk_model(opset_versions: List[Tuple[str, int]]) -> ModelProto: graph = helper.make_graph([], "my graph", [], []) return helper.make_model_gen_version( graph, opset_imports=[helper.make_opsetid(*pair) for pair in opset_versions], ) def test(opset_versions: List[Tuple[str, int]], ir_version: int) -> None: model = mk_model(opset_versions) self.assertEqual(model.ir_version, ir_version) # opset version 9 requires minimum ir_version 4 test([("", 9)], 4) test([("", 10)], 5) test([("", 11)], 6) test([("", 12)], 7) test([("", 13)], 7) test([("", 14)], 7) test([("", 15)], 8) test([("", 16)], 8) test([("", 17)], 8) test([("", 18)], 8) test([("", 19)], 9) test([("", 20)], 9) test([("", 21)], 10) # standard opset can be referred to using empty-string or "ai.onnx" test([("ai.onnx", 9)], 4) test([("ai.onnx.ml", 2)], 6) test([("ai.onnx.ml", 3)], 8) test([("ai.onnx.ml", 4)], 9) test([("ai.onnx.ml", 5)], 10) test([("ai.onnx.training", 1)], 7) # helper should pick *max* IR version required from all opsets specified. test([("", 10), ("ai.onnx.ml", 2)], 6) self.assertRaises(ValueError, mk_model, [("", 100)]) class TestHelperTensorFunctions(unittest.TestCase): def test_make_string_tensor(self) -> None: string_list = [s.encode("utf-8") for s in ["Amy", "Billy", "Cindy", "David"]] tensor = helper.make_tensor( name="test", data_type=TensorProto.STRING, dims=(2, 2), vals=string_list, raw=False, ) self.assertEqual(string_list, list(tensor.string_data)) def test_make_bfloat16_tensor(self) -> None: # numpy doesn't support bf16, so we have to compute the correct result manually np_array = np.array( [ [1.0, 2.0], [3.0, 4.0], [0.099853515625, 0.099365234375], [0.0998535081744, 0.1], [np.nan, np.inf], ], dtype=np.float32, ) np_results = np.array( [ [ struct.unpack("!f", bytes.fromhex("3F800000"))[0], # 1.0 struct.unpack("!f", bytes.fromhex("40000000"))[0], ], # 2.0 [ struct.unpack("!f", bytes.fromhex("40400000"))[0], # 3.0 struct.unpack("!f", bytes.fromhex("40800000"))[0], ], # 4.0 [ struct.unpack("!f", bytes.fromhex("3DCC0000"))[ 0 ], # round-to-nearest-even rounds down (0x8000) struct.unpack("!f", bytes.fromhex("3DCC0000"))[0], ], # round-to-nearest-even rounds up (0x8000) [ struct.unpack("!f", bytes.fromhex("3DCC0000"))[ 0 ], # round-to-nearest-even rounds down (0x7fff) struct.unpack("!f", bytes.fromhex("3DCD0000"))[0], ], # round-to-nearest-even rounds up (0xCCCD) [ struct.unpack("!f", bytes.fromhex("7FC00000"))[0], # NaN struct.unpack("!f", bytes.fromhex("7F800000"))[0], ], # inf ] ) tensor = helper.make_tensor( name="test", data_type=TensorProto.BFLOAT16, dims=np_array.shape, vals=np_array, ) self.assertEqual(tensor.name, "test") np.testing.assert_equal(np_results, numpy_helper.to_array(tensor)) def test_make_float8e4m3fn_tensor(self) -> None: y = helper.make_tensor( "zero_point", TensorProto.FLOAT8E4M3FN, [5], [0, 0.5, 1, 50000, 10.1] ) ynp = numpy_helper.to_array(y) expected = np.array([0, 0.5, 1, 448, 10], dtype=np.float32) np.testing.assert_equal(expected, ynp) def test_make_float8e4m3fnuz_tensor(self) -> None: y = helper.make_tensor( "zero_point", TensorProto.FLOAT8E4M3FNUZ, [7], [0, 0.5, 1, 50000, 10.1, -0.00001, 0.00001], ) ynp = numpy_helper.to_array(y) expected = np.array([0, 0.5, 1, 240, 10, 0, 0], dtype=np.float32) np.testing.assert_equal(expected, ynp) def test_make_float8e5m2_tensor(self) -> None: y = helper.make_tensor( "zero_point", TensorProto.FLOAT8E5M2, [5], [0, 0.5, 1, 50000, 96] ) ynp = numpy_helper.to_array(y) expected = np.array([0, 0.5, 1, 49152, 96], dtype=np.float32) np.testing.assert_equal(expected, ynp) def test_make_float8e5m2fnuz_tensor(self) -> None: y = helper.make_tensor( "zero_point", TensorProto.FLOAT8E5M2FNUZ, [7], [0, 0.5, 1, 50000, 96, -0.0000001, 0.0000001], ) ynp = numpy_helper.to_array(y) expected = np.array([0, 0.5, 1, 49152, 96, 0, 0], dtype=np.float32) np.testing.assert_equal(expected, ynp) def test_make_bfloat16_tensor_raw(self) -> None: # numpy doesn't support bf16, so we have to compute the correct result manually np_array = np.array( [ [1.0, 2.0], [3.0, 4.0], [0.099853515625, 0.099365234375], [0.0998535081744, 0.1], [np.nan, np.inf], ], dtype=np.float32, ) np_results = np.array( [ [ struct.unpack("!f", bytes.fromhex("3F800000"))[0], # 1.0 struct.unpack("!f", bytes.fromhex("40000000"))[0], ], # 2.0 [ struct.unpack("!f", bytes.fromhex("40400000"))[0], # 3.0 struct.unpack("!f", bytes.fromhex("40800000"))[0], ], # 4.0 [ struct.unpack("!f", bytes.fromhex("3DCC0000"))[0], # truncated struct.unpack("!f", bytes.fromhex("3DCB0000"))[0], ], # truncated [ struct.unpack("!f", bytes.fromhex("3DCC0000"))[0], # truncated struct.unpack("!f", bytes.fromhex("3DCC0000"))[0], ], # truncated [ struct.unpack("!f", bytes.fromhex("7FC00000"))[0], # NaN struct.unpack("!f", bytes.fromhex("7F800000"))[0], ], # inf ] ) # write out 16-bit of fp32 to create bf16 using truncation, no rounding def truncate(x): return x >> 16 values_as_ints = np_array.astype(np.float32).view(np.uint32).flatten() packed_values = truncate(values_as_ints).astype(np.uint16).tobytes() tensor = helper.make_tensor( name="test", data_type=TensorProto.BFLOAT16, dims=np_array.shape, vals=packed_values, raw=True, ) self.assertEqual(tensor.name, "test") np.testing.assert_equal(np_results, numpy_helper.to_array(tensor)) def test_make_float8e4m3fn_tensor_raw(self) -> None: expected = np.array([0, 0.5, 1, 448, 10], dtype=np.float32) f8 = np.array( [helper.float32_to_float8e4m3(x) for x in expected], dtype=np.uint8 ) packed_values = f8.tobytes() y = helper.make_tensor( name="test", data_type=TensorProto.FLOAT8E4M3FN, dims=list(expected.shape), vals=packed_values, raw=True, ) ynp = numpy_helper.to_array(y) np.testing.assert_equal(expected, ynp) def test_make_float8e4m3fnuz_tensor_raw(self) -> None: expected = np.array([0, 0.5, 1, 240, 10], dtype=np.float32) f8 = np.array( [helper.float32_to_float8e4m3(x, uz=True) for x in expected], dtype=np.uint8 ) packed_values = f8.tobytes() y = helper.make_tensor( name="test", data_type=TensorProto.FLOAT8E4M3FNUZ, dims=list(expected.shape), vals=packed_values, raw=True, ) ynp = numpy_helper.to_array(y) np.testing.assert_equal(expected, ynp) def test_make_float8e5m2_tensor_raw(self) -> None: expected = np.array([0, 0.5, 1, 49152, 10], dtype=np.float32) f8 = np.array( [helper.float32_to_float8e5m2(x) for x in expected], dtype=np.uint8 ) packed_values = f8.tobytes() y = helper.make_tensor( name="test", data_type=TensorProto.FLOAT8E5M2, dims=list(expected.shape), vals=packed_values, raw=True, ) ynp = numpy_helper.to_array(y) np.testing.assert_equal(expected, ynp) def test_make_float8e5m2fnuz_tensor_raw(self) -> None: expected = np.array([0, 0.5, 1, 49152, 10], dtype=np.float32) f8 = np.array( [helper.float32_to_float8e5m2(x, fn=True, uz=True) for x in expected], dtype=np.uint8, ) packed_values = f8.tobytes() y = helper.make_tensor( name="test", data_type=TensorProto.FLOAT8E5M2FNUZ, dims=list(expected.shape), vals=packed_values, raw=True, ) ynp = numpy_helper.to_array(y) np.testing.assert_equal(expected, ynp) @parameterized.parameterized.expand( itertools.product( (TensorProto.UINT4, TensorProto.INT4), ((5, 4, 6), (4, 6, 5), (3, 3), (1,), (2**10,)), ) ) @unittest.skipIf( version_utils.numpy_older_than("1.22.0"), "The test requires numpy 1.22.0 or later", ) def test_make_4bit_tensor(self, dtype, dims) -> None: type_range = { TensorProto.UINT4: (0, 15), TensorProto.INT4: (-8, 7), } data = np.random.randint( type_range[dtype][0], high=type_range[dtype][1] + 1, size=dims ) y = helper.make_tensor("y", dtype, data.shape, data) ynp = to_array_extended(y) np.testing.assert_equal(data, ynp) @parameterized.parameterized.expand( itertools.product( (TensorProto.UINT4, TensorProto.INT4), ((5, 4, 6), (4, 6, 5), (3, 3), (1,)) ) ) def test_make_4bit_raw_tensor(self, dtype, dims) -> None: type_range = { TensorProto.UINT4: (0, 15), TensorProto.INT4: (-8, 7), } data = np.random.randint( type_range[dtype][0], high=type_range[dtype][1] + 1, size=dims ) packed_data = helper.pack_float32_to_4bit( data, signed=(dtype == TensorProto.INT4) ) y = helper.make_tensor( "packed_int4", dtype, dims, packed_data.tobytes(), raw=True ) ynp = numpy_helper.to_array(y) np.testing.assert_equal(data, ynp) def test_make_sparse_tensor(self) -> None: values = [1.1, 2.2, 3.3, 4.4, 5.5] values_tensor = helper.make_tensor( name="test", data_type=TensorProto.FLOAT, dims=(5,), vals=values ) indices = [1, 3, 5, 7, 9] indices_tensor = helper.make_tensor( name="test_indices", data_type=TensorProto.INT64, dims=(5,), vals=indices ) dense_shape = [10] sparse = helper.make_sparse_tensor(values_tensor, indices_tensor, dense_shape) self.assertEqual(sparse.values, values_tensor) self.assertEqual(sparse.indices, indices_tensor) self.assertEqual(sparse.dims, dense_shape) def test_make_tensor_value_info(self) -> None: vi = helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 4)) checker.check_value_info(vi) # scalar value vi = helper.make_tensor_value_info("Y", TensorProto.FLOAT, ()) checker.check_value_info(vi) def test_make_sparse_tensor_value_info(self) -> None: vi = helper.make_sparse_tensor_value_info("X", TensorProto.FLOAT, (2, 3)) checker.check_value_info(vi) # scalar value vi = helper.make_sparse_tensor_value_info("Y", TensorProto.FLOAT, ()) checker.check_value_info(vi) class TestHelperOptionalAndSequenceFunctions(unittest.TestCase): def test_make_optional(self) -> None: values = [1.1, 2.2, 3.3, 4.4, 5.5] values_tensor = helper.make_tensor( name="test", data_type=TensorProto.FLOAT, dims=(5,), vals=values ) optional = helper.make_optional( name="test", elem_type=OptionalProto.TENSOR, value=values_tensor ) self.assertEqual(optional.name, "test") self.assertEqual(optional.elem_type, OptionalProto.TENSOR) self.assertEqual(optional.tensor_value, values_tensor) # Test Sequence values_sequence = helper.make_sequence( name="test", elem_type=SequenceProto.TENSOR, values=[values_tensor, values_tensor], ) optional = helper.make_optional( name="test", elem_type=OptionalProto.SEQUENCE, value=values_sequence ) self.assertEqual(optional.name, "test") self.assertEqual(optional.elem_type, OptionalProto.SEQUENCE) self.assertEqual(optional.sequence_value, values_sequence) # Test None optional_none = helper.make_optional( name="test", elem_type=OptionalProto.UNDEFINED, value=None ) self.assertEqual(optional_none.name, "test") self.assertEqual(optional_none.elem_type, OptionalProto.UNDEFINED) self.assertFalse(optional_none.HasField("tensor_value")) def test_make_optional_value_info(self) -> None: tensor_type_proto = helper.make_tensor_type_proto(elem_type=2, shape=[5]) tensor_val_into = helper.make_value_info( name="test", type_proto=tensor_type_proto ) optional_type_proto = helper.make_optional_type_proto(tensor_type_proto) optional_val_info = helper.make_value_info( name="test", type_proto=optional_type_proto ) self.assertEqual(optional_val_info.name, "test") self.assertTrue(optional_val_info.type.optional_type) self.assertEqual( optional_val_info.type.optional_type.elem_type, tensor_val_into.type ) # Test Sequence sequence_type_proto = helper.make_sequence_type_proto(tensor_type_proto) optional_type_proto = helper.make_optional_type_proto(sequence_type_proto) optional_val_info = helper.make_value_info( name="test", type_proto=optional_type_proto ) self.assertEqual(optional_val_info.name, "test") self.assertTrue(optional_val_info.type.optional_type) sequence_value_info = helper.make_value_info( name="test", type_proto=tensor_type_proto ) self.assertEqual( optional_val_info.type.optional_type.elem_type.sequence_type.elem_type, sequence_value_info.type, ) def test_make_seuence_value_info(self) -> None: tensor_type_proto = helper.make_tensor_type_proto(elem_type=2, shape=None) sequence_type_proto = helper.make_sequence_type_proto(tensor_type_proto) sequence_val_info = helper.make_value_info( name="test", type_proto=sequence_type_proto ) sequence_val_info_prim = helper.make_tensor_sequence_value_info( name="test", elem_type=2, shape=None ) self.assertEqual(sequence_val_info, sequence_val_info_prim) class TestPrintableGraph(unittest.TestCase): def test_initializer_with_matching_graph_input(self) -> None: add = helper.make_node("Add", ["X", "Y_Initializer"], ["Z"]) value_info = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1])] graph = helper.make_graph( [add], "test", [ helper.make_tensor_value_info("X", TensorProto.FLOAT, [1]), helper.make_tensor_value_info("Y_Initializer", TensorProto.FLOAT, [1]), ], # inputs [helper.make_tensor_value_info("Z", TensorProto.FLOAT, [1])], # outputs [ helper.make_tensor("Y_Initializer", TensorProto.FLOAT, [1], [1]) ], # initializers doc_string=None, value_info=value_info, ) graph_str = helper.printable_graph(graph) self.assertTrue( """) optional inputs with matching initializers ( %Y_Initializer[FLOAT, 1]""" in graph_str, graph_str, ) def test_initializer_no_matching_graph_input(self) -> None: add = helper.make_node("Add", ["X", "Y_Initializer"], ["Z"]) value_info = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1])] graph = helper.make_graph( [add], "test", [helper.make_tensor_value_info("X", TensorProto.FLOAT, [1])], # inputs [helper.make_tensor_value_info("Z", TensorProto.FLOAT, [1])], # outputs [ helper.make_tensor("Y_Initializer", TensorProto.FLOAT, [1], [1]) ], # initializers doc_string=None, value_info=value_info, ) graph_str = helper.printable_graph(graph) self.assertTrue( """) initializers ( %Y_Initializer[FLOAT, 1]""" in graph_str, graph_str, ) def test_unknown_dimensions(self) -> None: graph = helper.make_graph( [helper.make_node("Add", ["X", "Y_Initializer"], ["Z"])], "test", [helper.make_tensor_value_info("X", TensorProto.FLOAT, [None])], # inputs [helper.make_tensor_value_info("Z", TensorProto.FLOAT, [None])], # outputs [ helper.make_tensor("Y_Initializer", TensorProto.FLOAT, [1], [1]) ], # initializers doc_string=None, ) model = helper.make_model(graph) checker.check_model(model) graph_str = helper.printable_graph(graph) self.assertIn("X[FLOAT, ?]", graph_str) @pytest.mark.parametrize( "tensor_dtype", [ t for t in helper.get_all_tensor_dtypes() if t not in { TensorProto.BFLOAT16, TensorProto.FLOAT8E4M3FN, TensorProto.FLOAT8E4M3FNUZ, TensorProto.FLOAT8E5M2, TensorProto.FLOAT8E5M2FNUZ, TensorProto.UINT4, TensorProto.INT4, TensorProto.STRING, TensorProto.COMPLEX64, TensorProto.COMPLEX128, } ], ids=lambda tensor_dtype: helper.tensor_dtype_to_string(tensor_dtype), ) def test_make_tensor_vals(tensor_dtype: int) -> None: np_array = np.random.randn(2, 3).astype( helper.tensor_dtype_to_np_dtype(tensor_dtype) ) tensor = helper.make_tensor( name="test", data_type=tensor_dtype, dims=np_array.shape, vals=np_array ) np.testing.assert_equal(np_array, numpy_helper.to_array(tensor)) @pytest.mark.parametrize( "tensor_dtype", [ t for t in helper.get_all_tensor_dtypes() if t not in { TensorProto.BFLOAT16, TensorProto.STRING, TensorProto.FLOAT8E4M3FN, TensorProto.FLOAT8E4M3FNUZ, TensorProto.FLOAT8E5M2, TensorProto.FLOAT8E5M2FNUZ, TensorProto.UINT4, TensorProto.INT4, } ], ids=lambda tensor_dtype: helper.tensor_dtype_to_string(tensor_dtype), ) def test_make_tensor_raw(tensor_dtype: int) -> None: np_array = np.random.randn(2, 3).astype( helper.tensor_dtype_to_np_dtype(tensor_dtype) ) tensor = helper.make_tensor( name="test", data_type=tensor_dtype, dims=np_array.shape, vals=np_array.tobytes(), raw=True, ) np.testing.assert_equal(np_array, numpy_helper.to_array(tensor)) class TestHelperMappingFunctions(unittest.TestCase): # TODO (#4554): remove these tests about catching warnings after the deprecation period # Test these new functions should not raise any deprecation warnings @pytest.mark.filterwarnings("error::DeprecationWarning") def test_tensor_dtype_to_np_dtype_not_throw_warning(self) -> None: _ = helper.tensor_dtype_to_np_dtype(TensorProto.FLOAT) @pytest.mark.filterwarnings("error::DeprecationWarning") def test_tensor_dtype_to_storage_tensor_dtype_not_throw_warning(self) -> None: _ = helper.tensor_dtype_to_storage_tensor_dtype(TensorProto.FLOAT) @pytest.mark.filterwarnings("error::DeprecationWarning") def test_tensor_dtype_to_field_not_throw_warning(self) -> None: _ = helper.tensor_dtype_to_field(TensorProto.FLOAT) @pytest.mark.filterwarnings("error::DeprecationWarning") def test_np_dtype_to_tensor_dtype_not_throw_warning(self) -> None: _ = helper.np_dtype_to_tensor_dtype(np.dtype("float32")) def test_tensor_dtype_to_np_dtype_bfloat16(self) -> None: self.assertEqual( helper.tensor_dtype_to_np_dtype(TensorProto.BFLOAT16), np.dtype("float32") ) def test_tensor_dtype_to_storage_tensor_dtype_bfloat16(self) -> None: self.assertEqual( helper.tensor_dtype_to_storage_tensor_dtype(TensorProto.BFLOAT16), TensorProto.UINT16, ) # BFloat16 tensor uses TensorProto.UINT16 as storage type; # And the field name for TensorProto.UINT16 is int32_data def test_tensor_dtype_to_field_bfloat16(self) -> None: self.assertEqual( helper.tensor_dtype_to_field(TensorProto.BFLOAT16), "int32_data" ) class TestAttrTypeToStr(unittest.TestCase): @parameterized.parameterized.expand( [ (AttributeProto.AttributeType.FLOAT, "FLOAT"), # type: ignore[attr-defined] (AttributeProto.AttributeType.INT, "INT"), # type: ignore[attr-defined] (AttributeProto.AttributeType.STRING, "STRING"), # type: ignore[attr-defined] (AttributeProto.AttributeType.TENSOR, "TENSOR"), # type: ignore[attr-defined] (AttributeProto.AttributeType.GRAPH, "GRAPH"), # type: ignore[attr-defined] (AttributeProto.AttributeType.SPARSE_TENSOR, "SPARSE_TENSOR"), # type: ignore[attr-defined] (AttributeProto.AttributeType.TYPE_PROTO, "TYPE_PROTO"), # type: ignore[attr-defined] (AttributeProto.AttributeType.FLOATS, "FLOATS"), # type: ignore[attr-defined] (AttributeProto.AttributeType.INTS, "INTS"), # type: ignore[attr-defined] (AttributeProto.AttributeType.STRINGS, "STRINGS"), # type: ignore[attr-defined] (AttributeProto.AttributeType.TENSORS, "TENSORS"), # type: ignore[attr-defined] (AttributeProto.AttributeType.GRAPHS, "GRAPHS"), # type: ignore[attr-defined] (AttributeProto.AttributeType.SPARSE_TENSORS, "SPARSE_TENSORS"), # type: ignore[attr-defined] (AttributeProto.AttributeType.TYPE_PROTOS, "TYPE_PROTOS"), # type: ignore[attr-defined] ] ) def test_attr_type_to_str(self, attr_type, expected_str): result = helper._attr_type_to_str(attr_type) self.assertEqual(result, expected_str) def test_attr_type_to_str_undefined(self): result = helper._attr_type_to_str(9999) self.assertEqual(result, "UNDEFINED") if __name__ == "__main__": unittest.main() pytest.main([__file__])