diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index c093c7d3c..55e2e1d11 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -9146,8 +9146,12 @@ def multinomial(context, node): if num_samples > 1 and not replacement: raise ValueError("When num_samples is larger than 1, only replacement=True is supported.") # Based on PyTorch documentations, the input to `torch.multinomial` is probability, not logit. - x = mb.random_categorical(x=x, size=num_samples, mode="probs", name=node.name) - context.add(x) + samples = mb.random_categorical(x=x, size=num_samples, mode="probs") + # `torch.multinomial` returns int64-valued indices; `mb.random_categorical` keeps the + # input float dtype. Without this cast the converted model emits floats where downstream + # ops (e.g. `gather`/`index_select`) expect integer indices, matching #2337. + samples = mb.cast(x=samples, dtype="int32", name=node.name) + context.add(samples) @register_torch_op 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 4db09989a..d90d05d5d 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -14615,6 +14615,42 @@ def forward(self, x): # The counting of 1 in PyTorch and CoreML output should be similar. assert np.abs(np.sum(mlmodel_out) - np.sum(torch_out)) / mlmodel_out.size < 0.05 + def test_multinomial_returns_int(self): + """Regression test for https://github.com/apple/coremltools/issues/2337. + `torch.multinomial` returns an int64 tensor of indices; the converter + used to emit the raw float output of MIL `random_categorical`, which + broke downstream `gather` / `index_select` consumers. The converted + program must end with an int cast so the output dtype matches torch. + """ + + class TestModel(nn.Module): + def forward(self, x): + return torch.multinomial(x, num_samples=3, replacement=True) + + model = TestModel().eval() + input_data = torch.tensor([[0.1, 0.3, 0.5, 0.1], [0.5, 0.2, 0.2, 0.1]]) + torch_model = torch.jit.trace(model, input_data) + + prog = ct.convert( + torch_model, + inputs=[ct.TensorType(name="x", shape=input_data.shape)], + convert_to="milinternal", + minimum_deployment_target=ct.target.iOS17, + ) + + ops = get_op_types_in_program(prog) + assert "random_categorical" in ops + assert "cast" in ops, ( + "Expected an int cast after `random_categorical` so the output " + "dtype matches torch.multinomial's int64 contract (issue #2337)." + ) + + # The final output must be an integer dtype. + out = prog.functions["main"].outputs[0] + assert types.is_int(out.dtype), ( + f"multinomial output dtype must be int, got {out.dtype}" + ) + @pytest.mark.parametrize( "compute_unit, backend", itertools.product(compute_units, backends),