fuse_layernorm_or_instancenorm: check gamma / beta shapes, not just their rank

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

@@ -11,6 +11,7 @@
 from coremltools.converters.mil.mil import Block
 from coremltools.converters.mil.mil import Builder as mb
 from coremltools.converters.mil.mil import Operation, Program, Var
+from coremltools.converters.mil.mil.ops.defs.iOS15.normalization import layer_norm
 from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass
 from coremltools.converters.mil.mil.passes.helper import (
     _check_no_output_connection,
@@ -178,13 +179,26 @@ def _try_apply_transform(
         gamma_rank = gamma_var.rank if gamma_var is not None else -1
         beta_rank = beta_var.rank if beta_var is not None else -1
 
+        # layer_norm requires gamma and beta to have shape x.shape[axes] exactly. A gamma
+        # that only broadcasts over those dims has the right rank but the wrong shape, and
+        # cannot be folded into the op.
+        normalized_shape = [reduce_op.x.shape[a] for a in negative_axes]
+        has_normalized_shape_affine = (
+            gamma_var is not None
+            and beta_var is not None
+            and layer_norm._is_compatible_shape(list(gamma_var.shape), normalized_shape)
+            and layer_norm._is_compatible_shape(list(beta_var.shape), normalized_shape)
+        )
+
         if gamma_rank == len(axes) and beta_rank == len(axes):
             # axes for layer_norm must be [-1] or [-1, -2] or [-1, -2, -3] and so on
             if negative_axes == list(range(-len(negative_axes), 0)):
-                is_layernorm = True
+                is_layernorm = has_normalized_shape_affine
 
         if rank == 4 and negative_axes == [-3]:
-            is_layernorm = (gamma_var is None and beta_var is None) or (gamma_rank == 1 and beta_rank == 1)
+            is_layernorm = (gamma_var is None and beta_var is None) or (
+                gamma_rank == 1 and beta_rank == 1 and has_normalized_shape_affine
+            )
 
             if gamma_var:
                 ops_to_remove.append(gamma_var.op)
@@ -199,9 +213,13 @@ def _try_apply_transform(
                 beta_var = None
 
         if rank == 4 and (negative_axes == [-2, -1] or negative_axes == [-3, -2]):
+            # instance_norm's gamma and beta are per channel, so they must be C long.
+            # A rank 1 squeezed shape is not enough: a gamma broadcasting over a spatial
+            # dim also squeezes to rank 1, and fusing it moves it onto the channel axis.
+            channel = reduce_op.x.shape[-1 if negative_axes == [-3, -2] else -3]
             if (
-                len(np.squeeze(gamma_var.val).shape) == 1
-                and len(np.squeeze(beta_var.val).shape) == 1
+                np.squeeze(gamma_var.val).shape == (channel,)
+                and np.squeeze(beta_var.val).shape == (channel,)
             ):
                 is_instancenorm = True
             if negative_axes == [-3, -2]:

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

@@ -6098,6 +6098,53 @@ def prog(x):
             prog, {"x": shape}, expected_output_shapes={block.outputs[0].name: shape}
         )
 
+    @staticmethod
+    def _build_norm_pattern(shape, axes, gamma, beta):
+        """y = x * [gamma * rsqrt(var + eps)] + (beta - mean * [gamma * rsqrt(var + eps)])"""
+
+        @mb.program(input_specs=[mb.TensorSpec(shape=shape)])
+        def prog(x):
+            mean = mb.reduce_mean(x=x, axes=axes, keep_dims=True)
+            var = mb.reduce_mean(x=mb.square(x=mb.sub(x=x, y=mean)), axes=axes, keep_dims=True)
+            scale = mb.mul(x=mb.rsqrt(x=mb.add(x=var, y=1e-5)), y=gamma)
+            return mb.add(x=mb.mul(x=x, y=scale), y=mb.sub(x=beta, y=mb.mul(x=mean, y=scale)))
+
+        return prog
+
+    def test_layer_norm_gamma_only_broadcasts(self):
+        """
+        gamma has the rank layer_norm wants but not its shape: it only broadcasts over the
+        normalized axis. layer_norm requires gamma.shape == x.shape[axes], so this pattern
+        must be left alone rather than fused into an op that rejects it.
+        """
+        shape = (1, 2, 5)
+        prog = self._build_norm_pattern(
+            shape, [-1], np.float32([2.0]), np.float32([0.5])
+        )
+        apply_pass_and_basic_check(prog, "common::fuse_layernorm_or_instancenorm")
+        assert "layer_norm" not in get_op_types_in_program(prog)
+        assert_model_is_valid(prog, {"x": shape})
+
+    def test_instance_norm_gamma_on_a_spatial_axis(self):
+        """
+        gamma squeezes to rank 1 but broadcasts over H, not over the channel axis.
+        instance_norm's gamma is per channel, so fusing would move gamma onto a different
+        axis and give it the wrong length.
+        """
+        shape = (1, 3, 4, 5)
+        gamma = np.arange(1, 5, dtype=np.float32).reshape(1, 1, 4, 1)
+        prog = self._build_norm_pattern(shape, [-2, -1], gamma, np.zeros_like(gamma))
+        apply_pass_and_basic_check(prog, "common::fuse_layernorm_or_instancenorm")
+        assert "instance_norm" not in get_op_types_in_program(prog)
+        assert_model_is_valid(prog, {"x": shape})
+
+    def test_instance_norm_gamma_on_the_channel_axis_is_still_fused(self):
+        shape = (1, 3, 4, 5)
+        gamma = np.arange(1, 4, dtype=np.float32).reshape(1, 3, 1, 1)
+        prog = self._build_norm_pattern(shape, [-2, -1], gamma, np.zeros_like(gamma))
+        apply_pass_and_basic_check(prog, "common::fuse_layernorm_or_instancenorm")
+        assert get_op_types_in_program(prog) == ["instance_norm"]
+
     @pytest.mark.parametrize(
         "with_affine, constexpr_beta", itertools.product([True, False], [True, False])
     )