Fix the reflection padding optimization when slice_by_index's end_mask is not specified

coremltools/converters/mil/mil/passes/defs/optimize_tensor_operation.py

@@ -841,6 +841,20 @@ def apply(self, prog):
 
     @staticmethod
     def _match_pattern(concat_op, block):
+
+        def resolve_end_mask(slice_op):
+            # Helper function to resolve the end_mask attribute of the slice op if it is not present
+            if "end_mask" in slice_op.inputs:
+                return slice_op.inputs["end_mask"].val
+
+            if slice_op.inputs["end"].val is None:
+                return None
+
+            infered_mask = []
+            for end_idx, axis_dim in zip(slice_op.inputs["end"].val, slice_op.inputs["x"].shape):
+                infered_mask.append(end_idx == axis_dim)
+            return infered_mask
+
         if concat_op.op_type != "concat":
             return False
 
@@ -882,15 +896,11 @@ def _match_pattern(concat_op, block):
                 return False
 
             if end_mask is None:
-                end_mask = slice_op.inputs["end_mask"].val
+                end_mask = resolve_end_mask(slice_op)
+                if end_mask is None or False not in end_mask:
+                    return False
                 axis = list(end_mask).index(False, 0, len(end_mask))
 
-            if end_mask is None:
-                return False
-
-            if axis != list(end_mask).index(False, 0, len(end_mask)):
-                return False
-
             # Check that we're only taking a slice of size 1
             end = slice_op.inputs["end"].val
             begin = slice_op.inputs["begin"].val

coremltools/converters/mil/mil/passes/tests/test_passes.py

@@ -5308,6 +5308,53 @@ def prog(x):
             expected_output_shapes={block.outputs[0].name: (3, 2, 6, 8)},
         )
 
+    def test_success_no_end_mask(self):
+        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 6, 8))])
+        def prog(x1):
+            left = mb.slice_by_index(
+                x=x1, begin=[0, 0, 0, 1], end=[1, 2, 6, 2],
+            )
+            right = mb.slice_by_index(
+                x=x1, begin=[0, 0, 0, -2], end=[1, 2, 6, -1],
+            )
+            x = mb.concat(values=[left, x1, right], axis=3)
+
+            return x
+
+        prev_prog, _, block = apply_pass_and_basic_check(prog, "common::use_reflection_padding")
+        assert get_op_types_in_program(prev_prog) == ["slice_by_index", "slice_by_index", "concat"]
+        assert get_op_types_in_program(prog) == ["pad"]
+
+        inputs = {"x1": (1, 2, 6, 8)}
+        assert_model_is_valid(
+            prog,
+            inputs,
+            expected_output_shapes={block.outputs[0].name: (1, 2, 6, 10)},
+        )
+
+    def test_failure_all_true_mask(self):
+        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 6, 8))])
+        def prog(x1):
+            left = mb.slice_by_index(
+                x=x1, begin=[0, 0, 0, 1], end=[1, 2, 6, 8],
+            )
+            right = mb.slice_by_index(
+                x=x1, begin=[0, 0, 0, -2], end=[1, 2, 6, 7],
+            )
+            x = mb.concat(values=[left, x1, right], axis=3)
+
+            return x
+
+        prev_prog, _, block = apply_pass_and_basic_check(prog, "common::use_reflection_padding")
+        assert get_op_types_in_program(prev_prog) == ["slice_by_index", "slice_by_index", "concat"]
+        assert get_op_types_in_program(prog) == ["slice_by_index", "slice_by_index", "concat"]
+
+        inputs = {"x1": (1, 2, 6, 8)}
+        assert_model_is_valid(
+            prog,
+            inputs,
+            expected_output_shapes={block.outputs[0].name: (1, 2, 6, 16)},
+        )
 
 class TestDivideToMultiply:
     def test_divide_to_multiply(self):