Support a symbolic batch dim in torch narrow and group_norm

coremltools/converters/mil/frontend/torch/ops.py

@@ -778,14 +778,18 @@ def narrow(context, node):
     begin = [0] * len(x.shape)
     begin[dim.val] = start.val
 
-    end = list(x.shape)
+    # Only the narrowed dim has an end worth stating. Masking every other dim off
+    # keeps the slice independent of x.shape, which may hold symbols under a
+    # flexible input shape and cannot be baked into a constant.
+    end = [0] * len(x.shape)
+    end_mask = [True] * len(x.shape)
     end[dim.val] = start.val + length.val
+    end_mask[dim.val] = False
 
     # torch.narrow accepts a negative start, which counts from the end of the
     # dim. Such a slice reaches the end of the dim exactly when start + length
     # is 0, which slice_by_index would read as the absolute index 0 and turn
     # into an empty slice, so mask that end off instead.
-    end_mask = [False] * len(x.shape)
     if start.val < 0 and end[dim.val] == 0:
         end_mask[dim.val] = True
 
@@ -2851,15 +2855,21 @@ def instance_norm(context, node):
 def _group_norm_impl(x: Var, num_groups: int, weight: Var, bias: Var, eps: float) -> Var:
     n, c = x.shape[0], x.shape[1]  # at minimum (N, C) required
     num_groups = builtins.min(num_groups, c)
-    new_shape = [n, num_groups, c // num_groups]
     # optimization for non symbolic shapes. This get rids of 3 mil ops that required on dynamic shapes
-    if not any_symbolic(x.shape[2:]):
+    # Any symbolic dim has to be read from the input shape at run time, including the
+    # batch dim, which cannot be baked into a constant shape either.
+    if not any_symbolic(x.shape):
+        new_shape = [n, num_groups, c // num_groups]
         new_shape += [*x.shape[2:]]  # adds remaining dims
         input_shape = [*x.shape]  # n, c, *
     else:
         input_shape = mb.shape(x=x)
-        input_shape_sliced = mb.slice_by_size(x=input_shape, begin=[2], size=[-1])  # x_shape[2:]
-        new_shape = mb.concat(values=[new_shape, input_shape_sliced], axis=0)
+        batch_sliced = mb.slice_by_size(x=input_shape, begin=[0], size=[1])  # x_shape[:1]
+        new_shape_values = [batch_sliced, [num_groups, c // num_groups]]
+        if x.rank > 2:
+            # x_shape[2:]
+            new_shape_values.append(mb.slice_by_size(x=input_shape, begin=[2], size=[-1]))
+        new_shape = mb.concat(values=new_shape_values, axis=0)
 
     num_extra_axes = len(x.shape[2:])
     axes_ = [int(i) for i in range(2, 2 + num_extra_axes + 1)]

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -759,6 +759,33 @@ def forward(self, x):
                         compute_unit=compute_unit,
                     )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend",
+        itertools.product(compute_units, backends, frontends),
+    )
+    def test_narrow_dynamic_batch(self, compute_unit, backend, frontend):
+        """Narrowing one dim must not depend on the size of the other dims."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.narrow(x, 2, 1, 2)
+
+        lower_bound = 2
+        upper_bound_coreml = 10 if backend[0] == "mlprogram" else -1
+        upper_bound_torch = None if upper_bound_coreml == -1 else upper_bound_coreml
+        batch_coreml = RangeDim(default=4, lower_bound=lower_bound, upper_bound=upper_bound_coreml)
+        batch_torch = torch.export.Dim(name="batch", min=lower_bound, max=upper_bound_torch)
+
+        self.run_compare_torch(
+            (4, 3, 5),
+            Model(),
+            frontend=frontend,
+            backend=backend,
+            compute_unit=compute_unit,
+            converter_input_type=[TensorType(shape=(batch_coreml, 3, 5), dtype=np.float32)],
+            torch_export_dynamic_shapes={"x": {0: batch_torch}},
+        )
+
     @pytest.mark.parametrize(
         "compute_unit, backend, frontend, shape",
         itertools.product(
@@ -1510,6 +1537,40 @@ def test_groupnorm_dynamic(self, compute_unit, backend, frontend, group_features
             torch_export_dynamic_shapes=torch_export_dynamic_shapes,
         )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend, group_features, rank",
+        itertools.product(compute_units, backends, frontends, [(16, 32), (1, 1)], [2, 3, 4]),
+    )
+    def test_groupnorm_dynamic_batch(
+        self, compute_unit, backend, frontend, group_features, rank
+    ):
+        """Only the batch dim is dynamic; the rest of the shape is static."""
+        model = nn.GroupNorm(group_features[0], group_features[1])
+
+        lower_bound = 2
+        upper_bound_coreml = 10 if backend[0] == "mlprogram" else -1
+        upper_bound_torch = None if upper_bound_coreml == -1 else upper_bound_coreml
+        batch_coreml = RangeDim(default=6, lower_bound=lower_bound, upper_bound=upper_bound_coreml)
+        batch_torch = torch.export.Dim(name="batch", min=lower_bound, max=upper_bound_torch)
+
+        spatial_shape = (5,) * (rank - 2)
+        converter_input_type = [
+            TensorType(
+                shape=(batch_coreml, group_features[1]) + spatial_shape, dtype=np.float32
+            )
+        ]
+        torch_export_dynamic_shapes = {"input": {0: batch_torch}}
+
+        self.run_compare_torch(
+            (6, group_features[1]) + spatial_shape,
+            model,
+            frontend=frontend,
+            backend=backend,
+            compute_unit=compute_unit,
+            converter_input_type=converter_input_type,
+            torch_export_dynamic_shapes=torch_export_dynamic_shapes,
+        )
+
 
 class TestLinear(TorchBaseTest):
     @pytest.mark.parametrize(