diff --git a/examples/recipes/google_owlv2-large-patch14-ensemble/cpu/cpu/zero-shot-object-detection_fp16_config.json b/examples/recipes/google_owlv2-large-patch14-ensemble/cpu/cpu/zero-shot-object-detection_fp16_config.json new file mode 100644 index 000000000..e819bdfa0 --- /dev/null +++ b/examples/recipes/google_owlv2-large-patch14-ensemble/cpu/cpu/zero-shot-object-detection_fp16_config.json @@ -0,0 +1,97 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 1008, + 1008 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 49408 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + }, + { + "name": "pred_boxes" + }, + { + "name": "text_embeds" + }, + { + "name": "image_embeds" + } + ] + }, + "optim": {}, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "task": "zero-shot-object-detection", + "model_id": "google/owlv2-large-patch14-ensemble", + "model_type": "owlv2", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "zero-shot-object-detection", + "model_class": "AutoModelForZeroShotObjectDetection", + "model_type": "owlv2" + } +} \ No newline at end of file diff --git a/examples/recipes/google_owlv2-large-patch14-ensemble/cpu/cpu/zero-shot-object-detection_fp32_config.json b/examples/recipes/google_owlv2-large-patch14-ensemble/cpu/cpu/zero-shot-object-detection_fp32_config.json new file mode 100644 index 000000000..d2dbc34ca --- /dev/null +++ b/examples/recipes/google_owlv2-large-patch14-ensemble/cpu/cpu/zero-shot-object-detection_fp32_config.json @@ -0,0 +1,75 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 1008, + 1008 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 49408 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + }, + { + "name": "pred_boxes" + }, + { + "name": "text_embeds" + }, + { + "name": "image_embeds" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "zero-shot-object-detection", + "model_class": "AutoModelForZeroShotObjectDetection", + "model_type": "owlv2" + } +} \ No newline at end of file diff --git a/src/winml/modelkit/pattern/op_input_gen/__init__.py b/src/winml/modelkit/pattern/op_input_gen/__init__.py index 0ae982453..1f5ab62d7 100644 --- a/src/winml/modelkit/pattern/op_input_gen/__init__.py +++ b/src/winml/modelkit/pattern/op_input_gen/__init__.py @@ -6,6 +6,7 @@ from .binary_like_input_generator import * from .constant_of_shape_input_generator import ConstantOfShapeInputGenerator from .conv_input_generator import * +from .einsum_input_generator import EinsumInputGenerator from .expand_input_generator import ExpandInputGenerator from .flatten_input_generator import FlattenInputGenerator from .global_pooling_input_generator import * diff --git a/src/winml/modelkit/pattern/op_input_gen/einsum_input_generator.py b/src/winml/modelkit/pattern/op_input_gen/einsum_input_generator.py new file mode 100644 index 000000000..a8863d9fd --- /dev/null +++ b/src/winml/modelkit/pattern/op_input_gen/einsum_input_generator.py @@ -0,0 +1,220 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Input generator for the Einsum operator.""" + +from typing import Any + +from .op_input_gen import ( + InputShapeConstraint, + OpInputGenerator, + VariadicInputConstraint, + register_runtime_checker_op, +) + + +@register_runtime_checker_op +class EinsumInputGenerator(OpInputGenerator): + """Input generator for Einsum operator. + + Einsum signature: + - Inputs: *Inputs (variadic) - List of input tensors + - Attributes: equation (string, required) - The einsum equation string + - Output: Output - Result tensor + + Test coverage strategy: + - Common 2-input equations (matrix multiply, batched matmul, dot product, + outer product, element-wise multiply with reduction) + - Single-input equations (transpose, trace, diagonal) + - Various tensor ranks (1D through 4D) + """ + + op_name = "Einsum" + + def get_finite_attribute_sets(self) -> dict[str, list[Any]]: + """Return empty dict; equation is specified per combination.""" + return {} + + def get_input_and_infinite_attribute_combinations( + self, + ) -> list[dict[str, object]]: + """Return input combinations for Einsum. + + Each combination specifies the equation and matching input shapes. + Covers common real-world einsum patterns. + """ + combinations: list[dict[str, object]] = [ + # === Single-input equations === + # Transpose 2D: ij->ji + { + "Inputs": VariadicInputConstraint([InputShapeConstraint((3, 4))]), + "equation": "ij->ji", + }, + # Diagonal: ii->i + { + "Inputs": VariadicInputConstraint([InputShapeConstraint((4, 4))]), + "equation": "ii->i", + }, + # Sum all: ij-> + { + "Inputs": VariadicInputConstraint([InputShapeConstraint((3, 4))]), + "equation": "ij->", + }, + # === Two-input equations === + # Matrix multiply: ij,jk->ik + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((3, 4)), InputShapeConstraint((4, 5))] + ), + "equation": "ij,jk->ik", + }, + # Dot product: i,i-> + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((6,)), InputShapeConstraint((6,))] + ), + "equation": "i,i->", + }, + # Outer product: i,j->ij + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((3,)), InputShapeConstraint((4,))] + ), + "equation": "i,j->ij", + }, + # Element-wise multiply: ij,ij->ij + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((3, 4)), InputShapeConstraint((3, 4))] + ), + "equation": "ij,ij->ij", + }, + # Batched matrix multiply: bij,bjk->bik + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((2, 3, 4)), InputShapeConstraint((2, 4, 5))] + ), + "equation": "bij,bjk->bik", + }, + # Batched dot with ellipsis: ...ij,...jk->...ik + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((2, 3, 4)), InputShapeConstraint((2, 4, 5))] + ), + "equation": "...ij,...jk->...ik", + }, + # Inner product pattern from OWLv2: ...pd,...qd->...pq + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((2, 3, 4)), InputShapeConstraint((2, 5, 4))] + ), + "equation": "...pd,...qd->...pq", + }, + # 4D batched: abij,abjk->abik + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((2, 3, 4, 5)), InputShapeConstraint((2, 3, 5, 6))] + ), + "equation": "abij,abjk->abik", + }, + # Bilinear: ik,jk->ij (shared contraction dim) + { + "Inputs": VariadicInputConstraint( + [InputShapeConstraint((3, 4)), InputShapeConstraint((5, 4))] + ), + "equation": "ik,jk->ij", + }, + ] + + return combinations + + def derive_properties(self, properties: dict[str, Any]) -> dict[str, Any]: + """Derive additional properties for Einsum operator testing. + + Args: + properties: Base properties containing: + - Inputs_shape: tuple of shapes for each input tensor + - attr_equation: equation string + + Returns: + Updated properties with Einsum-specific derived values: + - num_inputs: number of input tensors + - Inputs_dim: max rank among input tensors + - has_ellipsis: whether the equation uses ellipsis notation + - has_explicit_output: whether the equation has '->' output spec + - is_full_reduction: whether output is scalar (empty after '->') + - has_repeated_labels: whether any input has repeated subscripts + - has_contraction: whether dimensions are contracted (summed out) + - output_num_labels: number of explicit output labels (output rank) + - inputs_share_all_labels: whether all inputs use identical label sets + """ + item = properties.copy() + + inputs_shape = item.get("Inputs_shape") + if inputs_shape is not None and len(inputs_shape) > 0: + item["num_inputs"] = len(inputs_shape) + # Max rank across all inputs + item["Inputs_dim"] = max( + (len(s) for s in inputs_shape if s is not None), default=0 + ) + else: + item["num_inputs"] = 0 + item["Inputs_dim"] = 0 + + # Derive semantic properties from the equation string + equation = item.get("attr_equation") + if equation is not None: + item["has_ellipsis"] = "..." in equation + item["has_explicit_output"] = "->" in equation + + # Parse output portion + output_part = equation.split("->")[1] if "->" in equation else "" + output_labels = set(output_part.replace("...", "")) + item["is_full_reduction"] = ( + "->" in equation and len(output_labels) == 0 + ) + + # Check for repeated labels in any single input term (e.g. "ii->i") + inputs_part = equation.split("->")[0] + input_terms = inputs_part.split(",") + item["has_repeated_labels"] = any( + len(t.replace("...", "")) != len(set(t.replace("...", ""))) + for t in input_terms + ) + + # Check for contraction (input labels absent from output) + all_input_labels = set( + "".join(t.replace("...", "") for t in input_terms) + ) + item["has_contraction"] = len(all_input_labels - output_labels) > 0 + + # Output rank (number of explicit output labels, excluding ellipsis) + # Analogous to Reshape's shape_len — EPs may differ by output dimensionality + item["output_num_labels"] = len(output_labels) + + # Whether all inputs share the same set of labels (broadcast/element-wise) + # vs having disjoint/partially-overlapping labels (contraction/outer product) + input_label_sets = [ + set(t.replace("...", "")) for t in input_terms + ] + if len(input_label_sets) >= 2: + item["inputs_share_all_labels"] = all( + s == input_label_sets[0] for s in input_label_sets[1:] + ) + else: + item["inputs_share_all_labels"] = True + + return item + + def get_infinite_property_names(self) -> list[str]: + """Returns names of infinite properties for Einsum operator. + + Input shapes are unbounded. The equation attribute is NOT included here + so that each distinct equation string participates in exact rule matching. + This is the conservative approach: rules will not generalize across + equations, avoiding false positive EP support claims for untested + equations with different contraction patterns, output permutations, + or broadcasting semantics. + """ + return ["Inputs_shape", "Inputs_value", "Inputs_is_constant"] diff --git a/tests/unit/analyze/test_input_generators.py b/tests/unit/analyze/test_input_generators.py index 613f867b3..a6e93517a 100644 --- a/tests/unit/analyze/test_input_generators.py +++ b/tests/unit/analyze/test_input_generators.py @@ -41,7 +41,7 @@ class TestInputGeneratorRegistry: def test_all_operators_registered(self) -> None: """Test that all operators are registered.""" # Verify count - assert len(get_registered_operators()) == 119 + assert len(get_registered_operators()) == 120 def test_get_runtime_checker_op(self) -> None: """Test retrieving operator generators by name.""" @@ -537,3 +537,255 @@ def test_conv_transpose_generated_cases_cover_derived_state_combinations( } assert has_odd_exact_stride_case assert has_even_non_exact_stride_case + + +class TestEinsumDeriveProperties: + """Test Einsum derive_properties extracts semantic features from equation.""" + + @pytest.fixture() + def einsum_generator(self): + """Create an Einsum generator for opset 22.""" + domain = ONNXDomain.AI_ONNX + schema = domain.get_op_schema("Einsum", 22) + generator_class = get_runtime_checker_op("Einsum") + return generator_class(schema) + + @pytest.mark.parametrize( + "equation,input_shapes,expected", + [ + # Matrix multiply: contraction on j + ( + "ij,jk->ik", + ((3, 4), (4, 5)), + { + "has_ellipsis": False, + "has_explicit_output": True, + "is_full_reduction": False, + "has_repeated_labels": False, + "has_contraction": True, + "output_num_labels": 2, + "inputs_share_all_labels": False, + }, + ), + # Bilinear: shared contraction dim k, different semantics from matmul + ( + "ik,jk->ij", + ((3, 4), (5, 4)), + { + "has_ellipsis": False, + "has_explicit_output": True, + "is_full_reduction": False, + "has_repeated_labels": False, + "has_contraction": True, + "output_num_labels": 2, + "inputs_share_all_labels": False, + }, + ), + # Batched matmul with ellipsis (OWLv2 pattern) + ( + "...pd,...qd->...pq", + ((2, 3, 4), (2, 5, 4)), + { + "has_ellipsis": True, + "has_explicit_output": True, + "is_full_reduction": False, + "has_repeated_labels": False, + "has_contraction": True, + "output_num_labels": 2, + "inputs_share_all_labels": False, + }, + ), + # Diagonal extraction: repeated labels + ( + "ii->i", + ((4, 4),), + { + "has_ellipsis": False, + "has_explicit_output": True, + "is_full_reduction": False, + "has_repeated_labels": True, + "has_contraction": False, + "output_num_labels": 1, + "inputs_share_all_labels": True, + }, + ), + # Full reduction: scalar output + ( + "ij->", + ((3, 4),), + { + "has_ellipsis": False, + "has_explicit_output": True, + "is_full_reduction": True, + "has_repeated_labels": False, + "has_contraction": True, + "output_num_labels": 0, + "inputs_share_all_labels": True, + }, + ), + # Outer product: no contraction + ( + "i,j->ij", + ((3,), (4,)), + { + "has_ellipsis": False, + "has_explicit_output": True, + "is_full_reduction": False, + "has_repeated_labels": False, + "has_contraction": False, + "output_num_labels": 2, + "inputs_share_all_labels": False, + }, + ), + # Element-wise multiply: no contraction, inputs share all labels + ( + "ij,ij->ij", + ((3, 4), (3, 4)), + { + "has_ellipsis": False, + "has_explicit_output": True, + "is_full_reduction": False, + "has_repeated_labels": False, + "has_contraction": False, + "output_num_labels": 2, + "inputs_share_all_labels": True, + }, + ), + # Batched matmul without ellipsis + ( + "bij,bjk->bik", + ((2, 3, 4), (2, 4, 5)), + { + "has_ellipsis": False, + "has_explicit_output": True, + "is_full_reduction": False, + "has_repeated_labels": False, + "has_contraction": True, + "output_num_labels": 3, + "inputs_share_all_labels": False, + }, + ), + ], + ) + def test_einsum_equation_semantic_properties( + self, einsum_generator, equation, input_shapes, expected + ): + """Verify each equation derives distinct semantic properties.""" + properties = { + "Inputs_shape": input_shapes, + "attr_equation": equation, + } + result = einsum_generator.derive_properties(properties) + + for key, value in expected.items(): + assert result[key] == value, ( + f"equation={equation!r}: expected {key}={value}, got {result[key]}" + ) + + def test_same_rank_different_equations_produce_different_properties( + self, einsum_generator + ): + """Equations with same arity and rank but different semantics must differ. + + This is the core regression test for the reviewer's concern: 'bij,bjk->bik' + and '...pd,...qd->...pq' have the same num_inputs=2 and Inputs_dim=3 but + must produce different derived property sets. + """ + props_batched = einsum_generator.derive_properties( + { + "Inputs_shape": ((2, 3, 4), (2, 4, 5)), + "attr_equation": "bij,bjk->bik", + } + ) + props_ellipsis = einsum_generator.derive_properties( + { + "Inputs_shape": ((2, 3, 4), (2, 5, 4)), + "attr_equation": "...pd,...qd->...pq", + } + ) + + # Same arity and rank + assert props_batched["num_inputs"] == props_ellipsis["num_inputs"] == 2 + assert props_batched["Inputs_dim"] == props_ellipsis["Inputs_dim"] == 3 + + # But different semantic properties (at minimum, has_ellipsis differs) + semantic_keys = [ + "has_ellipsis", + "has_explicit_output", + "is_full_reduction", + "has_repeated_labels", + "has_contraction", + ] + batched_sig = tuple(props_batched[k] for k in semantic_keys) + ellipsis_sig = tuple(props_ellipsis[k] for k in semantic_keys) + assert batched_sig != ellipsis_sig, ( + "Same-rank equations must produce different semantic signatures" + ) + + +class TestEinsumEquationExecution: + """Dedicated test that executes every Einsum equation independently. + + This addresses the reviewer's concern that validate_inputs() groups cases by + input shape constraints, causing same-shape-different-equation cases to be + skipped after the first in a group passes. This test builds a real ONNX Einsum + node for each equation and runs inference to verify the output shape. + """ + + @pytest.mark.parametrize( + "equation,input_shapes,expected_output_shape", + [ + # Single-input equations + ("ij->ji", [(3, 4)], (4, 3)), + ("ii->i", [(4, 4)], (4,)), + ("ij->", [(3, 4)], ()), + # Two-input equations + ("ij,jk->ik", [(3, 4), (4, 5)], (3, 5)), + ("i,i->", [(6,), (6,)], ()), + ("i,j->ij", [(3,), (4,)], (3, 4)), + ("ij,ij->ij", [(3, 4), (3, 4)], (3, 4)), + ("bij,bjk->bik", [(2, 3, 4), (2, 4, 5)], (2, 3, 5)), + ("...ij,...jk->...ik", [(2, 3, 4), (2, 4, 5)], (2, 3, 5)), + ("...pd,...qd->...pq", [(2, 3, 4), (2, 5, 4)], (2, 3, 5)), + ("abij,abjk->abik", [(2, 3, 4, 5), (2, 3, 5, 6)], (2, 3, 4, 6)), + ("ik,jk->ij", [(3, 4), (5, 4)], (3, 5)), + ], + ) + def test_einsum_equation_output_shape( + self, equation, input_shapes, expected_output_shape + ): + """Execute each equation via ORT and verify output shape independently.""" + import onnxruntime as ort + from onnx import TensorProto, helper + from onnx.shape_inference import infer_shapes + + # Build a minimal ONNX model with a single Einsum node + inputs = [] + input_tensors = [] + for i, shape in enumerate(input_shapes): + name = f"input_{i}" + inputs.append(helper.make_tensor_value_info(name, TensorProto.FLOAT, shape)) + input_tensors.append(np.random.randn(*shape).astype(np.float32)) + + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, None) + + node = helper.make_node( + "Einsum", + inputs=[f"input_{i}" for i in range(len(input_shapes))], + outputs=["output"], + equation=equation, + ) + + graph = helper.make_graph([node], "test_einsum", inputs, [output]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model = infer_shapes(model) + + # Run inference + sess = ort.InferenceSession(model.SerializeToString()) + feed = {f"input_{i}": t for i, t in enumerate(input_tensors)} + result = sess.run(None, feed) + + assert result[0].shape == expected_output_shape, ( + f"equation={equation!r}: expected shape {expected_output_shape}, " + f"got {result[0].shape}" + )