# Copyright (c) ONNX Project Contributors # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import locale import platform import pytest import onnx from onnx import GraphProto, OperatorSetIdProto, TensorProto, checker class TestBasicFunctions: def check_graph(self, graph: GraphProto) -> None: assert len(graph.node) == 3 assert graph.node[0].op_type == "MatMul" assert graph.node[1].op_type == "Add" assert graph.node[2].op_type == "Softmax" def test_parse_graph(self) -> None: input = """ agraph (float[N, 128] X, float[128,10] W, float[10] B) => (float[N] C) { T = MatMul(X, W) S = Add(T, B) C = Softmax(S) } """ graph = onnx.parser.parse_graph(input) self.check_graph(graph) def test_parse_model(self) -> None: input = """ < ir_version: 7, opset_import: [ "" : 10, "com.microsoft": 1] > agraph (float[N, 128] X, float[128,10] W, float[10] B) => (float[N] C) { T = MatMul(X, W) S = Add(T, B) C = Softmax(S) } """ model = onnx.parser.parse_model(input) assert model.ir_version == 7 assert len(model.opset_import) == 2 self.check_graph(model.graph) def test_parse_graph_error(self) -> None: input = """ agraph (float[N, 128] X, float[128,10] W, float[10] B) => (float[N] C) { T = MatMul[X, W] S = Add(T, B) C = Softmax(S) } """ with pytest.raises(onnx.parser.ParseError): onnx.parser.parse_graph(input) def test_parse_model_error(self) -> None: input = """ < ir_version: 7, opset_import: [ "" : 10 "com.microsoft": 1] > agraph (float[N, 128] X, float[128,10] W, float[10] B) => (float[N] C) { T = MatMul(X, W) S = Add(T, B) C = Softmax(S) } """ with pytest.raises(onnx.parser.ParseError): onnx.parser.parse_model(input) def test_parse_function_with_attributes(self) -> None: input = """ < ir_version: 9, opset_import: [ "" : 15, "custom_domain" : 1], producer_name: "FunctionProtoTest", producer_version: "1.0", model_version: 1, doc_string: "A test model for model local functions." > agraph (float[N] x) => (float[N] out) { out = custom_domain.Selu(x) } < domain: "custom_domain", opset_import: [ "" : 15], doc_string: "Test function proto" > Selu (X) => (C) { constant_alpha = Constant() constant_gamma = Constant() alpha_x = CastLike(constant_alpha, X) gamma_x = CastLike(constant_gamma, X) exp_x = Exp(X) alpha_x_exp_x = Mul(alpha_x, exp_x) alpha_x_exp_x_ = Sub(alpha_x_exp_x, alpha_x) neg = Mul(gamma_x, alpha_x_exp_x_) pos = Mul(gamma_x, X) _zero = Constant() zero = CastLike(_zero, X) less_eq = LessOrEqual(X, zero) C = Where(less_eq, neg, pos) } """ model = onnx.parser.parse_model(input) checker.check_model(model) @pytest.mark.parametrize( "graph_text, expected_attribute", [ ( "agraph (float[N] x) => (float[N] out) { out = custom_domain.Selu(x) }", {}, ), ( "agraph (float[N] x) => (float[N] out) { out = custom_domain.Selu(x) }", {"alpha": 2.0}, ), ( "agraph (float[N] x) => (float[N] out) { out = custom_domain.Selu(x) }", {"gamma": 3.0}, ), ( "agraph (float[N] x) => (float[N] out) { out = custom_domain.Selu(x) }", {"alpha": 2.0, "gamma": 3.0}, ), ], ) def test_composite_parse_function_with_attributes( self, graph_text: str, expected_attribute: dict ) -> None: default_alpha = 1.67326319217681884765625 default_gamma = 1.05070102214813232421875 def expect_custom_node_attribute(node, attributes): for key in attributes: match_attr = [attr for attr in node.attribute if attr.name == key] assert len(match_attr) == 1 assert match_attr[0].f == attributes[key] def expect_model_function_attribute(model): assert len(model.functions[0].attribute_proto) == 2 attr_proto_alpha = [ attr_proto for attr_proto in model.functions[0].attribute_proto if attr_proto.name == "alpha" ] assert len(attr_proto_alpha) == 1 and attr_proto_alpha[0].f == default_alpha attr_proto_gamma = [ attr_proto for attr_proto in model.functions[0].attribute_proto if attr_proto.name == "gamma" ] assert len(attr_proto_gamma) == 1 and attr_proto_gamma[0].f == default_gamma function_text = f""" < domain: "custom_domain", opset_import: [ "" : 15], doc_string: "Test function proto" > Selu (X) => (C) {{ constant_alpha = Constant() constant_gamma = Constant() alpha_x = CastLike(constant_alpha, X) gamma_x = CastLike(constant_gamma, X) exp_x = Exp(X) alpha_x_exp_x = Mul(alpha_x, exp_x) alpha_x_exp_x_ = Sub(alpha_x_exp_x, alpha_x) neg = Mul(gamma_x, alpha_x_exp_x_) pos = Mul(gamma_x, X) _zero = Constant() zero = CastLike(_zero, X) less_eq = LessOrEqual(X, zero) C = Where(less_eq, neg, pos) }} """ functions = [onnx.parser.parse_function(function_text)] graph = onnx.parser.parse_graph(graph_text) opset_imports = [ OperatorSetIdProto(domain="", version=15), OperatorSetIdProto(domain="custom_domain", version=1), ] model = onnx.helper.make_model( graph, functions=functions, opset_imports=opset_imports ) checker.check_model(model) expect_model_function_attribute(model) expect_custom_node_attribute(model.graph.node[0], expected_attribute) def test_parse_node(self): node = onnx.parser.parse_node( "out1, out2 = SomeDomain.SomeOp (in1, in2)" ) assert list(node.input) == ["in1", "in2"] assert list(node.output) == ["out1", "out2"] assert len(node.attribute) == 1 attr_val = onnx.helper.get_node_attr_value(node, "attr1") assert attr_val == 1 assert node.domain == "SomeDomain" assert node.op_type == "SomeOp" def test_parse_float_attribute_from_int_literal(self): model = onnx.parser.parse_model( """ < ir_version: 9, opset_import: [ "" : 18, "custom_domain" : 1] > agraph (float[N] x) => (float[N] out) { out = custom_domain.Foo(x) } """ ) attr = model.graph.node[0].attribute[0] assert attr.type == onnx.AttributeProto.FLOAT assert attr.HasField("f") assert not attr.HasField("i") assert attr.f == 2.0 def test_missing_identifier(self): node = onnx.parser.parse_node("= SomeOp ()") assert list(node.input) == [] assert list(node.output) == [] node = onnx.parser.parse_node(", = SomeOp (,)") assert list(node.input) == [""] assert list(node.output) == [""] node = onnx.parser.parse_node("x, = SomeOp (y,)") assert list(node.input) == ["y"] assert list(node.output) == ["x"] node = onnx.parser.parse_node(",x = SomeOp (,y)") assert list(node.input) == ["", "y"] assert list(node.output) == ["", "x"] def test_quoted_empty_identifier(self): node = onnx.parser.parse_node('"" = SomeOp ("")') assert list(node.input) == [""] assert list(node.output) == [""] node = onnx.parser.parse_node('"",x = SomeOp ("",y)') assert list(node.input) == ["", "y"] assert list(node.output) == ["", "x"] def test_quoted_string_symbolic_dim(self): # Test parsing a quoted string as a symbolic dimension (non-identifier dim_param) graph = onnx.parser.parse_graph( 'agraph (float["M + N"] x) => (float["M + N"] y) { y = Identity(x) }' ) assert graph.input[0].type.tensor_type.shape.dim[0].dim_param == "M + N" assert graph.output[0].type.tensor_type.shape.dim[0].dim_param == "M + N" @pytest.mark.parametrize( "test_literal, expect_exception", [ ("not_a_good_float", True), ("inf1", True), ("-inf1", True), ("nan0", True), ("-nan0", True), ("naninf", True), ("inf", False), ("-inf", False), ("infinity", False), ("-infinity", False), ("nan", False), ("-NaN", False), ], ) def test_parse_various_float_values(self, test_literal, expect_exception): model_text = f""" < ir_version: 8, opset_import: ["" : 18, "this" : 1], producer_name: "FunctionProtoTest", producer_version: "1.0" > _func () => () {{ tmp = Constant () }} """ if expect_exception: with pytest.raises(onnx.parser.ParseError): onnx.parser.parse_model(model_text) else: model = onnx.parser.parse_model(model_text) assert model.ir_version == 8 assert model.producer_name == "FunctionProtoTest" assert model.producer_version == "1.0" assert len(model.graph.node) == 1 assert len(model.graph.node[0].attribute) == 1 assert model.graph.node[0].attribute[0].name == "value_float" assert model.graph.node[0].attribute[0].type == onnx.AttributeProto.FLOAT assert str(model.graph.node[0].attribute[0].f) == str(float(test_literal)) @pytest.mark.parametrize( "name, itype", [ ("bfloat16", TensorProto.BFLOAT16), ("bool", TensorProto.BOOL), ("complex64", TensorProto.COMPLEX64), ("complex128", TensorProto.COMPLEX128), ("double", TensorProto.DOUBLE), ("float16", TensorProto.FLOAT16), ("float", TensorProto.FLOAT), ("float8e4m3fn", TensorProto.FLOAT8E4M3FN), ("float8e4m3fnuz", TensorProto.FLOAT8E4M3FNUZ), ("float8e5m2", TensorProto.FLOAT8E5M2), ("float8e5m2fnuz", TensorProto.FLOAT8E5M2FNUZ), ("int2", TensorProto.INT2), ("int4", TensorProto.INT4), ("int8", TensorProto.INT8), ("int16", TensorProto.INT16), ("int32", TensorProto.INT32), ("int64", TensorProto.INT64), ("string", TensorProto.STRING), ("uint2", TensorProto.UINT2), ("uint4", TensorProto.UINT4), ("uint8", TensorProto.UINT8), ("uint16", TensorProto.UINT16), ("uint32", TensorProto.UINT32), ("uint64", TensorProto.UINT64), ("float4e2m1", TensorProto.FLOAT4E2M1), ], ) def test_parse_graph_types(self, name, itype) -> None: w = '{"0"}' if itype == TensorProto.STRING else "{0}" text_graph = f""" < ir_version: 10, opset_import: [ "" : 19] > agraph (float[N] X) => ({name}[N] C) < {name}[1] weight = {w} > {{ C = Cast(X) }} """ graph = onnx.parser.parse_model(text_graph) assert len(graph.graph.node) == 1 def test_locale_independent_float_parsing(self) -> None: """Regression test: float parsing must work under non-US locales. See https://github.com/onnx/onnx/issues/8111 """ original_locale = locale.setlocale(locale.LC_NUMERIC, None) def restore_locale() -> None: locale.setlocale(locale.LC_NUMERIC, original_locale) # Try to set a locale with comma as decimal separator. # Use platform-appropriate locale names. is_windows = platform.system() == "Windows" candidates = ( ("German_Germany.1252", "French_France.1252") if is_windows else ("de_DE.UTF-8", "fr_FR.UTF-8") ) locale_set = False for candidate in candidates: try: locale.setlocale(locale.LC_NUMERIC, candidate) locale_set = True break except locale.Error: continue if not locale_set: restore_locale() pytest.skip("No locale with comma decimal separator available") try: model_text = """ agraph (float[1, 5] X) => (float[1, 5] Y) { Y = LeakyRelu (X) } """ model = onnx.parser.parse_model(model_text) node = model.graph.node[0] assert node.attribute[0].name == "alpha" alpha = node.attribute[0].f assert alpha == pytest.approx(0.123, abs=1e-5), ( "Float attribute misparsed under non-US locale" ) finally: restore_locale()