NotImplementedError: Only tensor assignment with exactly 1 pure dimension selection is supported

[JierunChen]

🌱 Describe your Feature Request

Coremltools 5.2 only supports tensor assignment with exactly 1 pure dimension selection. I'm trying to convert a model from pytorch to coreml, which involves more flexible tensor assignments, as in the use case. After running the code, it raises a NotImplementedError("Only tensor assignment with exactly 1 pure dimension selection is supported"). It would be great if the feature could be supported.

Use case

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


class Model(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        x[:, :4, :-1, :-1] = 0
        return x


if __name__ == "__main__":
    x = torch.randn(1, 3, 8, 8)
    model = Model()
    model.eval()
    y = model(x)

    torchscript_model = torch.jit.trace(model, x)
    coreml_model = ct.convert(
        torchscript_model,
        inputs=[ct.ImageType(name="x_1",
                             shape=x.shape)]
    )

[TobyRoseman]

I don't understand this feature request. Can you give us a simple standalone model which could only be converted by adding this feature request?

[JierunChen]

@TobyRoseman Hi, I've updated the feature request and posted a standalone use case. Thanks!

[TobyRoseman]

Thanks @JierunChen - I can reproduce the issue (in coremltools 6.0b1).

[HashedViking]

@JierunChen Hello, do you have a workaround for this problem?

[JierunChen]

@JierunChen Hello, do you have a workaround for this problem?

I temporally use torch.cat to concatenate a series of fragments. But it probably would be slower.

[den-run-ai]

same issue here, torch.cat worked great!