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(