I looked more into this and this error rises only when using models with states. For models without states, model prediction does not raise this error even if the model combines several inputs with flexible and static shapes.
Here is a code snippet to reproduce
import traceback
import numpy as np
import coremltools as ct
import torch
import torch.nn as nn
class ModelA(nn.Module):
def __init__(self):
super().__init__()
d = torch.zeros(1, dtype=torch.float16)
self.register_buffer("d", d)
def forward(self, a, b, c):
return a + b + c + self.d
class ModelB(nn.Module):
def __init__(self):
super().__init__()
d = torch.zeros(1, dtype=torch.float16)
self.register_buffer("d", d)
def forward(self, a, b, c):
c0 = c[0]
return a + b + c0 + self.d
def _convert(model, states, c_is_enum):
with torch.inference_mode():
traced_model = torch.jit.trace(
model,
(
torch.randn(1, dtype=torch.float16),
torch.randn(1, dtype=torch.float16),
torch.randn(1, dtype=torch.float16),
),
)
inputs = [
ct.TensorType(
name="a",
shape=ct.EnumeratedShapes([(1,), (2,)]),
dtype=np.float16,
),
ct.TensorType(
name="b",
shape=ct.EnumeratedShapes([(1,), (2,)]),
dtype=np.float16,
),
ct.TensorType(
name="c",
shape=ct.EnumeratedShapes([(1,), (2,)]) if c_is_enum else (1,),
dtype=np.float16,
),
]
outputs = [
ct.TensorType(
name="output",
dtype=np.float16,
),
]
mlmodel = ct.convert(
traced_model,
convert_to="milinternal",
inputs=inputs,
outputs=outputs,
states=states,
minimum_deployment_target=ct.target.iOS18,
compute_units=ct.ComputeUnit.CPU_AND_NE,
compute_precision=ct.precision.FLOAT16,
)
print(mlmodel)
mlmodel = ct.convert(
mlmodel,
inputs=inputs,
outputs=outputs,
minimum_deployment_target=ct.target.iOS18,
compute_units=ct.ComputeUnit.CPU_AND_NE,
compute_precision=ct.precision.FLOAT16,
)
return mlmodel
def convert(model, with_state, c_is_enum):
torch.random.manual_seed(42)
np.random.seed(42)
if with_state:
states = [
ct.StateType(
wrapped_type=ct.TensorType(shape=(1,)),
name="d",
),
]
mlmodel = _convert(model, states, c_is_enum=c_is_enum)
state = mlmodel.make_state()
else:
mlmodel = _convert(model, None, c_is_enum=c_is_enum)
state = None
a = torch.randn(2, dtype=torch.float16)
b = torch.randn(2, dtype=torch.float16)
c = torch.randn(1, dtype=torch.float16)
if c_is_enum:
c = c.repeat(2)
print(a, b, c)
print(
mlmodel.predict(
{
"a": a,
"b": b,
"c": c,
},
state,
)
)
torch.random.manual_seed(42)
np.random.seed(42)
model_a = ModelA().eval()
model_b = ModelB().eval()
try:
convert(model_a, with_state=True, c_is_enum=False)
except Exception as e:
traceback.print_exc()
try:
convert(model_a, with_state=False, c_is_enum=False)
except Exception as e:
traceback.print_exc()
try:
convert(model_b, with_state=True, c_is_enum=True)
except Exception as e:
traceback.print_exc()
In the code I perform 3 conversions.
- The first one that combines flexible, static shapes and states raises the error
libc++abi: terminating due to uncaught exception of type CoreML::MLNeuralNetworkUtilities::AsymmetricalEnumeratedShapesException:
A model doesn't allow input features with enumerated flexibility to have unequal number of enumerated shapes, but input feature a has 2 enumerated shapes and input feature c has 1 enumerated shapes.
- The second conversion that does not use state and combines flexible and static shapes works correctly.
- The last conversion in which in which I convert my static input to have enumerated shape, and tile the tensor to repeat the value along all the dimensions, but index to use only the first value inside of the model, runs prediction without raising the error.
