Torch state buffer `copy_` fails to convert
- 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
Assessment
This issue has not been assessed yet.