apple / apple/coremltools

Torch state buffer `copy_` fails to convert

Open
#2,272 1 comment 0 reactions 1 assignee Claimed by @jakesabathia2 View on GitHub
bug
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug
I'm trying to use the new State feature and I want to overwrite the entire state with a new value. I'm able to write a MIL model that does this, but when I try to use `copy_` to do this from PyTorch, I get a conversion error.

## Stack Trace
```
Traceback (most recent call last):
File "/torch_convert.py", line 52, in
model = ct.convert(torch.jit.trace(Net().eval(), x),
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lib/python3.11/site-packages/coremltools/converters/_converters_entry.py", line 635, in convert
mlmodel = mil_convert(
^^^^^^^^^^^^
File "/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 188, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
proto, mil_program = mil_convert_to_proto(
^^^^^^^^^^^^^^^^^^^^^
File "/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 288, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 108, in __call__
return load(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 74, in load
converter = TorchConverter(
^^^^^^^^^^^^^^^
File "/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 582, in __init__
p(self.graph)
File "/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/torchir_passes.py", line 153, in generate_tensor_assignment_ops
raise ValueError("No matching select or slice.")
ValueError: No matching select or slice.
```

## To Reproduce
```python
import coremltools as ct
from coremltools.converters.mil.mil import Builder as mb, types

import numpy as np
import torch
from torch import nn

# MIL
@mb.program(input_specs=[mb.TensorSpec((10,10), dtype=types.fp16),
mb.StateTensorSpec((10,10), dtype=types.fp16),],
opset_version=ct.target.iOS18,
)
def prog(x, buf_state):
buf_value = mb.read_state(input=buf_state)
x_sq = mb.mul(x=x, y=x)
mb.coreml_update_state(state=buf_state, value=x_sq)

y = mb.add(x=buf_value, y=x_sq)
return y

prog_model = ct.convert(prog,minimum_deployment_target=ct.target.iOS18, compute_units=ct.ComputeUnit.CPU_ONLY)
prog_model.save("prog.mlpackage")

print("MIL prog outputs:")
state = prog_model.make_state()
print(prog_model.predict({"x": torch.ones(10,10)*1}, state))
print(prog_model.predict({"x": torch.ones(10,10)*2}, state))

# Torch
class Net(nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("buf", torch.zeros(10,10).half())

def forward(self, x):
buf_value = self.buf.clone()
x_sq = x * x

# coremltools/converters/mil/frontend/torch/torchir_passes.py", line 153, in generate_tensor_assignment_ops
# raise ValueError("No matching select or slice.")
self.buf.copy_(x_sq)

return x_sq + buf_value

net = Net().eval()
print("torch outputs:")
print(net(torch.ones(10,10)*1))
print(net(torch.ones(10,10)*2))

x = torch.randn(10,10)
model = ct.convert(torch.jit.trace(Net().eval(), x),
inputs=[ct.TensorType(shape=x.shape, dtype=np.float16)],
states=[ct.StateType(wrapped_type=ct.TensorType(shape=(10,10)), name="buf")],
minimum_deployment_target=ct.target.iOS18)

model.save("torchmodel.mlpackage")
```

## System environment (please complete the following information):
- coremltools version: 8.0b1
- OS (e.g. MacOS version or Linux type): Version 15.0 Beta (24A5264n) (beta 1)
- Any other relevant version information (e.g. PyTorch or TensorFlow version): torch==2.3.0

## Additional context
If there are other ways to overwrite the entire state, happy to try them.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.