Incorrect behavior for torch.linspace

[twoertwein]

🐞Describing the bug

When end in linspace(start=constant, end=dynamic, steps=constant) is dependent on an array of unknown length, linspace behaves unexpected in two ways:

• CoreMLTools can't infer the shape of the output from linspace (even though steps is a constant)
• The number of elements of the output is sometimes larger than expected (should always have exactly steps elements)

Both issues do not exist when end is a constant.

To Reproduce

import coremltools
import torch


class Test(torch.nn.Module):
    def forward(self, x):
        # create an array of unknown length
        x = torch.nonzero(x >= 0.5).view(-1)
        # linspace should always have steps elements, but CoreML tools can't infer the fixed shape
        # and CoreMLTools is not compatible with pytorch:
        # steps=10 -> 11 elements (incorrect)
        # steps=5 -> 5 elements (correct)
        return torch.linspace(0, x.numel() - 1, steps=10)

data = torch.arange(0, 1000, dtype=torch.float32)
model = torch.jit.trace(Test(), example_kwarg_inputs={"x": data})

core_model = coremltools.convert(
    model,
    convert_to="mlprogram",
    inputs=[coremltools.TensorType(shape=data.shape, name="x")],
    outputs=[coremltools.TensorType(name="test")],
)

print(core_model)  # model shape missing for output!

actual = core_model.predict({"x": data.numpy()})["test"]
expected = model(data).numpy()

print(expected)
print(actual)

System environment (please complete the following information):

• coremltools version: 8.1
• OS (e.g. MacOS version or Linux type): Darwin Kernel Version 24.1.0 (M2)
• Any other relevant version information (e.g. PyTorch or TensorFlow version): pytorch 2.5.1

[twoertwein]

PyTorch's implementation basically does (+dtypes handling)

arange = torch.arange(0, steps)
step =  (end - start) / (steps - 1)
return arange * step + start

[TobyRoseman]

Using PyTorch version 2.4.0 (the max version we currently support). I get the following for actual:

array([  0.     , 110.86182, 221.72363, 332.58545, 443.44727, 554.3091 ,
       665.1709 , 776.0327 , 886.89453, 997.75635, 998.     ],
      dtype=float32)

And the following for expected:

array([  0.      , 110.888885, 221.77777 , 332.66666 , 443.55554 ,
       554.44446 , 665.3334  , 776.2222  , 887.1111  , 998.      ],
      dtype=float32)

Looks like actual has an extra value at [-2]