Accept an unflatten size read from the input shape

[LeSingh1]

Problem

torch.unflatten with a size read off the input aborts conversion once that dimension is flexible:

import torch, torch.nn as nn, coremltools as ct

class M(nn.Module):
    def forward(self, x):
        return torch.unflatten(x, 0, (x.shape[0], 1))

ct.convert(
    torch.jit.trace(M().eval(), torch.rand(3, 4)),
    inputs=[ct.TensorType(shape=(ct.RangeDim(lower_bound=1, upper_bound=10), 4))],
    convert_to="mlprogram",
)

File "coremltools/converters/mil/mil/builder.py", line 122, in _add_const
    raise ValueError("Cannot add const {}".format(val))
ValueError: Cannot add const [<Var object at 0x...>, <Var object at 0x...>]

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 —

target_shape = mb.concat(values=(pre_shape, unflattened_size_var, post_shape), axis=0)

— but hands that python list straight to mb.concat, which tries to turn it into a const and fails.

Fix

Stack the list elements into the rank 1 shape tensor concat expects. A size that is already const is untouched: it still binds to a single Var and takes the existing path, so nothing about the static case changes.

Test

TestUnflatten::test_unflatten_size_read_from_input_shape unflattens a flexible leading dimension into (n, 1) and (1, n) with n read from x.shape, and compares against torch. Both splits keep the product equal to the split dimension for every size the RangeDim admits, so the model is valid across the whole range rather than only at the sizes the test happens to feed.

All 4 parametrizations fail with Cannot add const [...] without the fix and pass with it.

The existing test_unflatten cannot reach this path: it builds nn.Unflatten from a python list fixed at construction time, so its unflattened size is always a compile-time constant even in its dynamic=True cases.

TorchScript only — torch.export resolves the size symbolically instead of emitting a list, so it never produces this binding.

Related

There is a long-open PR #2050 that attacks the same symptom from the other end, by making _array_construct itself emit a concat when a ListConstruct depends on graph inputs. This change is deliberately narrower — it only teaches unflatten to accept the list form the current _array_construct already produces, and would simply become a no-op if #2050 were to land.

Verification

pytest coremltools/converters/mil/frontend/torch/test/test_torch_ops.py -k TestUnflatten

passes on macOS / Apple silicon, torch 2.12, with the pre-existing xfails/skips unchanged.