Accept an unflatten size read from the input shape

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

@@ -887,6 +887,17 @@ def unflatten(context, node):
     if dim < 0:
         dim += x.rank
 
+    if isinstance(unflattened_size_var, (list, tuple)):
+        # A size read off the input, e.g. x.unflatten(1, (2, x.shape[1] // 2)), comes
+        # from a prim::ListConstruct whose elements are only known at run time, so it
+        # binds to a python list of scalar Vars instead of to a single const Var.
+        # Stack them into the rank 1 shape tensor that concat needs.
+        unflattened_size_var = mb.concat(
+            values=[mb.expand_dims(x=mb.cast(x=size, dtype="int32"), axes=[0])
+                    for size in unflattened_size_var],
+            axis=0,
+        )
+
     x_shape = mb.shape(x=x)
     pre_shape = mb.slice_by_index(x=x_shape, begin=[0], end=[dim])
     post_shape = mb.slice_by_index(x=x_shape, begin=[dim + 1], end=[len(x.shape)])

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

@@ -6291,6 +6291,41 @@ def forward(self, x):
             compute_unit=compute_unit,
         )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, size_position",
+        itertools.product(compute_units, backends, ["first", "last"]),
+    )
+    def test_unflatten_size_read_from_input_shape(self, compute_unit, backend, size_position):
+        """
+        ``torch.unflatten(x, 0, (x.shape[0], 1))`` reads a size off the input, so the
+        unflattened size comes from a ``prim::ListConstruct`` of run-time values and
+        binds to a python list of Vars rather than to a single const Var.
+
+        ``test_unflatten`` above cannot reach this: it builds ``nn.Unflatten`` from a
+        python list fixed at construction time, so its sizes are always const.
+
+        Splitting into ``(n, 1)`` / ``(1, n)`` keeps the product equal to the split
+        dimension for every size the RangeDim admits.
+        """
+
+        class Model(nn.Module):
+            def forward(self, x):
+                if size_position == "first":
+                    return torch.unflatten(x, 0, (x.shape[0], 1))
+                return torch.unflatten(x, 0, (1, x.shape[0]))
+
+        self.run_compare_torch(
+            (3, 4),
+            Model().eval(),
+            # torch.export resolves the size symbolically instead of emitting a list
+            frontend=TorchFrontend.TORCHSCRIPT,
+            backend=backend,
+            compute_unit=compute_unit,
+            converter_input_type=[
+                ct.TensorType(shape=(ct.RangeDim(lower_bound=1, upper_bound=10), 4))
+            ],
+        )
+
 
 class TestGather(TorchBaseTest):
     @pytest.mark.parametrize(