From 79994561a27610d8ddaaec0885026dad401df3bf Mon Sep 17 00:00:00 2001 From: LeSingh1 Date: Sat, 8 Aug 2026 19:18:52 -0700 Subject: [PATCH] Accept an unflatten size read from the input shape torch.unflatten(x, 0, (x.shape[0], 1)) over a flexible dimension aborts conversion with ValueError: Cannot add const [, ] A size read off the input comes from a prim::ListConstruct whose elements are only known at run time, so _get_bindings binds it to a python list of scalar Vars rather than to a single const Var. The converter already assembles the target shape dynamically, but handed that python list straight to mb.concat, which tried to turn it into a const. Stack the elements into the rank 1 shape tensor concat expects. A size that is already const is unaffected: it still binds to a single Var and takes the existing path. test_unflatten could not reach this because it builds nn.Unflatten from a python list fixed at construction time, so its sizes are always const. --- .../converters/mil/frontend/torch/ops.py | 11 ++++++ .../mil/frontend/torch/test/test_torch_ops.py | 35 +++++++++++++++++++ 2 files changed, 46 insertions(+) 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(