diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index 7f747baf4..b0d5f89cb 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/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)]) diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index 58860798b..2241884a8 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/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(