From c06cc0169bf3faf228b3c5f11286689791121bbe Mon Sep 17 00:00:00 2001 From: LeSingh1 Date: Sun, 9 Aug 2026 15:06:15 -0700 Subject: [PATCH] State the limitation when chunk / unbind meet a dynamic dim Both ops derive their number of outputs from the size of the dim they act on, so that dim has to be known at conversion time. Neither said so. constantchunk did float(x.shape[dim]) and unbind did [1] * x.shape[dim], which on a symbolic dim raise "TypeError: Cannot convert expression to float" and "TypeError: can't multiply sequence by non-int of type 'Symbol'" from inside the converter, with nothing naming the op or the limitation. Raise a ValueError that says which dim is the problem and what to do instead, as _adaptive_pool2d already does for the same situation. No model that converts today is affected: these paths only ever raised. Also adds the first chunk tests; there were none. --- .../converters/mil/frontend/torch/ops.py | 14 ++++++ .../mil/frontend/torch/test/test_torch_ops.py | 47 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index 7f747baf4..05f2e39c1 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -6263,6 +6263,13 @@ def _parse_keyword_args(context, node, dim) -> Var: if isinstance(dim, Var): dim = dim.val + if is_symbolic(x.shape[dim]): + raise ValueError( + f"unbind on dim {dim} of node {node.name} needs that dim to be known at " + "conversion time, because it decides how many outputs there are, but it is " + "dynamic. Use a fixed size for that dim." + ) + split_sizes = [1] * x.shape[dim] if len(split_sizes) == 1: res = [mb.squeeze(x=x, axes=[dim])] @@ -6373,6 +6380,13 @@ def constantchunk(context, node): dim = node.attr["dim"] total = x.shape[dim] + if is_symbolic(total): + raise ValueError( + f"chunk on dim {dim} of node {node.name} needs that dim to be known at " + "conversion time, because it decides how many chunks there are, but it is " + "dynamic. Use a fixed size for that dim, or torch.split with an explicit " + "split size." + ) size = int(_math.ceil(float(total) / float(chunks))) split_sizes = [size] * int(_math.floor(total / size)) remainder = total - sum(split_sizes) 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..5a4c08550 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -7752,6 +7752,53 @@ def test_unbind_one_dim_shape(self, compute_unit, backend, frontend): input_shape, model, compute_unit=compute_unit, backend=backend, frontend=frontend ) + @staticmethod + def test_unbind_dynamic_dim_is_rejected(): + """The unbound dim decides how many outputs there are, so it cannot be dynamic.""" + + class Model(nn.Module): + def forward(self, x): + return torch.cat(torch.unbind(x, dim=0), dim=0) + + traced = torch.jit.trace(Model().eval(), torch.rand(3, 4)) + with pytest.raises(ValueError, match="how many outputs there are"): + ct.convert( + traced, + inputs=[ct.TensorType(name="x", shape=(ct.RangeDim(2, 8), 4))], + minimum_deployment_target=ct.target.iOS17, + ) + + +class TestChunk(TorchBaseTest): + @pytest.mark.parametrize( + "compute_unit, backend, frontend, chunks", + itertools.product(compute_units, backends, frontends, [2, 3]), + ) + def test_chunk(self, compute_unit, backend, frontend, chunks): + class Model(nn.Module): + def forward(self, x): + return torch.cat(torch.chunk(x, chunks, dim=0), dim=0) + + self.run_compare_torch( + (7, 4), Model(), compute_unit=compute_unit, backend=backend, frontend=frontend + ) + + @staticmethod + def test_chunk_dynamic_dim_is_rejected(): + """The chunked dim decides how many chunks there are, so it cannot be dynamic.""" + + class Model(nn.Module): + def forward(self, x): + return torch.cat(torch.chunk(x, 3, dim=0), dim=0) + + traced = torch.jit.trace(Model().eval(), torch.rand(7, 4)) + with pytest.raises(ValueError, match="how many chunks there are"): + ct.convert( + traced, + inputs=[ct.TensorType(name="x", shape=(ct.RangeDim(2, 16), 4))], + minimum_deployment_target=ct.target.iOS17, + ) + class TestTranspose(TorchBaseTest): @pytest.mark.parametrize(