Spaces:
Sleeping
Sleeping
| # SPDX-License-Identifier: Apache-2.0 | |
| # Copyright (c) ONNX Project Contributors | |
| import unittest | |
| from typing import Dict, List, Optional, Tuple, Union | |
| import numpy as np | |
| import onnx | |
| from onnx import TensorProto, TypeProto | |
| from onnx.checker import ValidationError | |
| from onnx.defs import OpSchema, get_all_schemas_with_history, get_schema | |
| from onnx.helper import ( | |
| make_graph, | |
| make_node, | |
| make_opsetid, | |
| make_tensor_type_proto, | |
| make_tensor_value_info, | |
| ) | |
| from onnx.numpy_helper import from_array | |
| from onnx.shape_inference import InferenceError, infer_node_outputs | |
| ADD_SCHEMA = max( | |
| (s for s in get_all_schemas_with_history() if s.name == "Add" and s.domain == ""), | |
| key=lambda s: s.since_version, | |
| ) | |
| RESHAPE_SCHEMA = max( | |
| ( | |
| s | |
| for s in get_all_schemas_with_history() | |
| if s.name == "Reshape" and s.domain == "" | |
| ), | |
| key=lambda s: s.since_version, | |
| ) | |
| def _to_tensor_types( | |
| tensor_types: Dict[str, Tuple[int, Tuple[Union[int, str, None], ...]]] | |
| ) -> Dict[str, TypeProto]: | |
| return {key: make_tensor_type_proto(*value) for key, value in tensor_types.items()} | |
| def _run_case( | |
| schema: OpSchema, | |
| input_names: List[str], | |
| output_names: List[str], | |
| input_types: Dict[str, TypeProto], | |
| input_data: Optional[Dict[str, np.ndarray]] = None, | |
| ) -> Dict[str, TypeProto]: | |
| if input_data is None: | |
| input_data = {} | |
| return infer_node_outputs( | |
| schema, | |
| make_node(schema.name, input_names, output_names, domain=schema.domain), | |
| input_types, | |
| {key: from_array(arr) for key, arr in input_data.items()}, | |
| ) | |
| class TestInferenceFunctionCall(unittest.TestCase): | |
| def test_add_inference(self) -> None: | |
| cases = [ | |
| ( | |
| {"A": (TensorProto.FLOAT, ()), "B": (TensorProto.FLOAT, ())}, | |
| {"C": (TensorProto.FLOAT, ())}, | |
| ), | |
| ( | |
| { | |
| "A": (TensorProto.FLOAT, (None, 2)), | |
| "B": (TensorProto.FLOAT, (2,)), | |
| }, | |
| {"C": (TensorProto.FLOAT, (None, 2))}, | |
| ), | |
| ( | |
| { | |
| "A": (TensorProto.FLOAT, (None, 2)), | |
| "B": (TensorProto.FLOAT, (1, 2)), | |
| }, | |
| {"C": (TensorProto.FLOAT, (None, 2))}, | |
| ), | |
| ( | |
| { | |
| "A": (TensorProto.DOUBLE, ("n", "m")), | |
| "B": (TensorProto.DOUBLE, (1, "n", "m")), | |
| }, | |
| {"C": (TensorProto.DOUBLE, (1, "n", "m"))}, | |
| ), | |
| ( | |
| { | |
| "A": (TensorProto.FLOAT, ("x", 2)), | |
| "B": (TensorProto.FLOAT, ("y", 2)), | |
| }, | |
| {"C": (TensorProto.FLOAT, (None, 2))}, | |
| ), | |
| ] | |
| for ins, outs in cases: | |
| assert _run_case(ADD_SCHEMA, ["A", "B"], ["C"], _to_tensor_types(ins)) == _to_tensor_types(outs) # type: ignore | |
| def test_add_inference_raises_errors(self) -> None: | |
| with self.assertRaises(ValidationError): | |
| _run_case( | |
| ADD_SCHEMA, | |
| ["A"], | |
| ["C"], | |
| _to_tensor_types({"A": (TensorProto.FLOAT, (3, 4))}), | |
| ) | |
| with self.assertRaises(ValidationError): | |
| _run_case( | |
| ADD_SCHEMA, | |
| ["A", "B"], | |
| ["C"], | |
| _to_tensor_types({"A": (TensorProto.FLOAT, (3, 4)), "B": (2, (3, 4))}), | |
| ) | |
| with self.assertRaises(InferenceError): | |
| _run_case( | |
| ADD_SCHEMA, | |
| ["A", "B"], | |
| ["C"], | |
| _to_tensor_types( | |
| { | |
| "A": (TensorProto.FLOAT, (2, 4)), | |
| "B": (TensorProto.FLOAT, (3, 4)), | |
| } | |
| ), | |
| ) | |
| with self.assertRaises(KeyError): | |
| _run_case( | |
| ADD_SCHEMA, | |
| ["A", "B"], | |
| ["C"], | |
| _to_tensor_types({"A": (TensorProto.FLOAT, (3, 4))}), | |
| ) | |
| def test_reshape_inference(self) -> None: | |
| assert _run_case( | |
| RESHAPE_SCHEMA, | |
| ["x", "t"], | |
| ["y"], | |
| _to_tensor_types( | |
| { | |
| "x": (TensorProto.FLOAT, (5, 4)), | |
| "t": (TensorProto.INT64, (3,)), | |
| } | |
| ), | |
| {"t": np.array([2, 2, 5], dtype=np.int64)}, | |
| ) == _to_tensor_types({"y": (TensorProto.FLOAT, (2, 2, 5))}) | |
| def test_scan_inference_with_subgraph(self) -> None: | |
| seq_len = "sequence" | |
| input_size = 2 | |
| loop_state_size = 3 | |
| input_value_infos = [ | |
| make_tensor_value_info("loop_state_in", TensorProto.UNDEFINED, None), | |
| make_tensor_value_info("input", TensorProto.UNDEFINED, None), | |
| make_tensor_value_info("outer", TensorProto.UNDEFINED, None), | |
| ] | |
| output_value_infos = [ | |
| make_tensor_value_info("loop_state_out", TensorProto.UNDEFINED, None), | |
| make_tensor_value_info("output", TensorProto.FLOAT, (seq_len, input_size)), | |
| ] | |
| subgraph = make_graph( | |
| [ | |
| make_node("Identity", ["loop_state_in"], ["loop_state_out"]), | |
| make_node("Add", ["input", "outer"], ["output"]), | |
| ], | |
| "subgraph", | |
| input_value_infos, | |
| output_value_infos, | |
| ) | |
| assert infer_node_outputs( | |
| get_schema("Scan", 9), | |
| make_node( | |
| "Scan", | |
| ["loop_state_orig", "scan_input", "scan_outer"], | |
| ["loop_state_final", "scan_output"], | |
| num_scan_inputs=1, | |
| body=subgraph, | |
| ), | |
| _to_tensor_types( | |
| { | |
| "loop_state_orig": (TensorProto.FLOAT, (loop_state_size,)), | |
| "scan_input": (TensorProto.FLOAT, (seq_len, input_size)), | |
| "scan_outer": (TensorProto.FLOAT, (input_size,)), | |
| } | |
| ), | |
| # Same as default value in Scan-9 | |
| opset_imports=[make_opsetid("", 9)], | |
| ir_version=4, | |
| ) == _to_tensor_types( | |
| { | |
| "loop_state_final": (TensorProto.FLOAT, (loop_state_size,)), | |
| "scan_output": (TensorProto.FLOAT, (seq_len, input_size)), | |
| } | |
| ) | |
| def test_inference_with_conflow(self) -> None: | |
| model_script = """ | |
| < | |
| ir_version: 8, | |
| opset_import: ["" : 18, "onnxscript.atenlib" : 1], | |
| producer_name: "pytorch", | |
| producer_version: "2.1.0" | |
| > | |
| torch_jit (float input_0) => (float reault, int64 index) | |
| { | |
| reault, index = onnxscript.atenlib.aten_min_dim <dim = 0, keepdim = 1> (input_0) | |
| } | |
| < | |
| domain: "onnxscript.atenlib", | |
| opset_import: ["" : 18] | |
| > | |
| aten_min_dim <dim>(self) => (result_7, indices_6) | |
| { | |
| tmp = Shape (self) | |
| tmp_0 = Size (tmp) | |
| tmp_1 = Constant <value = int64 tmp_1 {0}> () | |
| tmp_1_cast = CastLike (tmp_1, tmp_0) | |
| tmp_2 = Equal (tmp_0, tmp_1_cast) | |
| cond = Not (tmp_2) | |
| indices_6, result_7 = If (cond) < | |
| then_branch = thenGraph_4 () => ( indices, result) { | |
| dim = Constant <value_int: int = @dim> () | |
| tmp_3 = Constant <value_ints = [-1]> () | |
| dims = Reshape (dim, tmp_3) | |
| result = ReduceMin <keepdims: int = @keepdim> (self, dims) | |
| indices = ArgMin <axis: int = @dim, keepdims: int = @keepdim> (self) | |
| }, else_branch = elseGraph_4 () => ( indices_4, result_5) { | |
| indices_4 = Constant <value_int = 0> () | |
| result_5 = Identity (self) | |
| } | |
| > | |
| } | |
| """ | |
| model = onnx.parser.parse_model(model_script) | |
| onnx.shape_inference.infer_shapes(model, strict_mode=False) | |
| with self.assertRaises(onnx.shape_inference.InferenceError): | |
| onnx.shape_inference.infer_shapes(model, strict_mode=True) | |
| def test_inference_with_attribute(self) -> None: | |
| model_script = """ | |
| < | |
| ir_version: 8, | |
| opset_import: ["" : 18, "custom" : 1], | |
| producer_name: "", | |
| producer_version: "1.0" | |
| > | |
| MeanVarianceNormalization (float[N] x) => (float[M] y) | |
| { | |
| y = custom.custom_mvn <axes = [0]> (x) | |
| } | |
| < | |
| domain: "custom", | |
| opset_import: ["" : 18] | |
| > | |
| custom_mvn <axes>(X) => (Y) | |
| { | |
| Exponent = Constant <value = float {2.0}>() | |
| Epsilon = Constant <value = float {1e-9}>() | |
| axes = Constant <value_ints: ints = @axes>() | |
| X_RM = ReduceMean (X, axes) | |
| EX_squared = Pow (X_RM, Exponent) | |
| X_squared = Pow (X, Exponent) | |
| E_Xsquared = ReduceMean (X_squared, axes) | |
| Variance = Sub (E_Xsquared, EX_squared) | |
| STD = Sqrt (Variance) | |
| X_variance = Sub (X, X_RM) | |
| Processed_STD = Add (STD, Epsilon) | |
| Y = Div (X_variance, Processed_STD) | |
| } | |
| """ | |
| model = onnx.parser.parse_model(model_script) | |
| # onnx.shape_inference.infer_shapes(model, strict_mode=False) | |
| onnx.shape_inference.infer_shapes(model, strict_mode=True) | |
| if __name__ == "__main__": | |
| unittest.main() | |