# Bug Report: EnumeratedShapes / FlexibleShapeInfo models fail at inference, no workaround found

## Describing the bug

A model compiled with `ct.EnumeratedShapes` cannot be used for inference from any available Apple API. The model **converts and saves without error**, and the `.mlpackage` is structurally valid. However, every input construction method fails at prediction time with the same underlying strides conflict from the Espresso runtime (`E5RT`).

We tested every documented and undocumented path in both Python and Swift — all fail identically. `MLShapedArray` was specifically tested as a potential workaround and also fails.

**Affected APIs:**
- Python: `coremltools model.predict()`
- Swift: `MLMultiArray(shape:dataType:)`
- Swift: `MLShapedArray` + `MLMultiArray(shaped:)`
- Swift: `MLDictionaryFeatureProvider` with `MLFeatureValue(multiArray:)`

---

## Stack Trace

**Python** — silent failure (batch dimension collapsed to 1, no exception):
```
# Input shape: (8, 256) → output shape: (1, 256)
# Batch dimension silently collapsed, no error raised
```

**Swift — both MLMultiArray and MLShapedArray:**
```
Type of context in function main's I/O contains unknown strides.
Using unknown strides for MIL tensor buffers with unknown shapes is not
recommended in E5ML. Please use row_alignment_in_bytes property instead.

E5RT encountered an STL exception.
msg = tensor_buffer has known strides while the model has FlexibleShapeInfo.
Strides must be unknown on all dimensions.

E5RT: tensor_buffer has known strides while the model has FlexibleShapeInfo.
Strides must be unknown on all dimensions. (11)

[Espresso::handle_ex_plan] exception=Espresso exception: "Invalid state":
Stack_nd layer: Invalid shapes of input tensors. status=-5

Unable to compute the prediction using a neural network model. It can be an
invalid input data or broken/unsupported model (error code: -5).
```

---

## To Reproduce

The reproduction is split into two steps: **conversion** (succeeds) and **inference** (fails). Both steps are required to reproduce the full issue.

### Step 1 — Convert model (succeeds without error)

```python
import coremltools as ct
import numpy as np
import torch
import torch.nn as nn

# Minimal model
class SimpleModel(nn.Module):
def forward(self, x): # x: [B, 256]
return x * 2.0

model = SimpleModel().eval()
traced = torch.jit.trace(model, torch.zeros(1, 256))

# Compile with EnumeratedShapes — conversion succeeds
ct_input = ct.TensorType(
name="context",
shape=ct.EnumeratedShapes(
shapes=[[1, 256], [4, 256], [8, 256], [16, 256], [32, 256]],
default=[1, 256],
),
dtype=np.float32,
)

mlmodel = ct.convert(
traced,
inputs=[ct_input],
minimum_deployment_target=ct.target.macOS13,
compute_units=ct.ComputeUnit.CPU_ONLY,
)

# Saves without error — model is structurally valid
mlmodel.save("test_enum.mlpackage")
print("Conversion: OK")
```

### Step 2 — Python inference (fails silently)

```python
import coremltools as ct
import numpy as np

# Load the saved model
loaded = ct.models.MLModel("test_enum.mlpackage")

# Batch size 1 — silent failure: batch dim collapsed to 1
out1 = loaded.predict({"context": np.ones((1, 256), dtype=np.float32)})
print("batch=1 shape:", np.array(list(out1.values())[0]).shape)
# Expected: (1, 256) — happens to work by coincidence (default shape)

# Batch size 8 — silent failure: output is (1, 256) not (8, 256)
out8 = loaded.predict({"context": np.ones((8, 256), dtype=np.float32)})
print("batch=8 shape:", np.array(list(out8.values())[0]).shape)
# Expected: (8, 256)
# Actual: (1, 256) ← batch dimension silently collapsed to default
```

### Step 3 — Swift inference (hard failure, all input construction methods)

```swift
import CoreML
import Foundation

// Load the model
let modelURL = Bundle.main.url(forResource: "test_enum", withExtension: "mlpackage")!
let config = MLModelConfiguration()
config.computeUnits = .cpuOnly

guard let model = try? MLModel(contentsOf: modelURL, configuration: config) else {
fatalError("Failed to load model")
}

// --- Attempt 1: MLMultiArray(shape:dataType:) — FAILS ---
do {
let arr = try MLMultiArray(shape: [1, 256], dataType: .float32)
// fill with test data
for i in 0..<256 { arr[i] = NSNumber(value: Float(i) / 256.0) }

let input = try MLDictionaryFeatureProvider(
dictionary: ["context": MLFeatureValue(multiArray: arr)]
)
let _ = try model.prediction(from: input)
} catch {
print("Attempt 1 failed:", error)
// → E5RT: tensor_buffer has known strides while model has FlexibleShapeInfo
}

// --- Attempt 2: MLShapedArray → MLMultiArray — ALSO FAILS ---
// MLShapedArray was tested as a potential workaround. It is NOT a workaround.
do {
let scalars = (0..<256).map { Float($0) / 256.0 }
let shaped = MLShapedArray(scalars: scalars, shape: [1, 256])
let arr = MLMultiArray(shaped)

let input = try MLDictionaryFeatureProvider(
dictionary: ["context": MLFeatureValue(multiArray: arr)]
)
let _ = try model.prediction(from: input)
} catch {
print("Attempt 2 failed:", error)
// → E5RT: tensor_buffer has known strides while model has FlexibleShapeInfo
// Identical error — changing the input construction API does not help.
}
```

---

## System environment

- **coremltools version:** 9.0
- **OS:** macOS 26.5 (M2 MacBook Pro 32 GB)
- **Xcode:** 26.5
- **PyTorch:** 2.11.0
- **Python:** 3.12.13
- **Device tested:** M2 MacBook Pro (macOS), iPhone 16 Pro (iOS 26), Apple Watch SE2 (watchOS 11)

---

## Additional context

### Root cause hypothesis

`EnumeratedShapes` models compiled to `mlprogram` format embed `FlexibleShapeInfo` metadata in the MIL graph. At prediction time, the Espresso runtime (`E5RT`) requires that all input tensors have *unknown* strides meaning it expects to control the memory layout of the input buffer itself. There is no public API to create an `MLMultiArray` with explicitly unknown strides. Every available construction path (`MLMultiArray(shape:)`, `MLShapedArray`, `bytesNoCopy:`) commits a concrete memory layout, which E5RT rejects.

### Suggested fix

Concretely: when `E5RT` receives an input for a model with `FlexibleShapeInfo`, if the input has known contiguous strides and the shape is in the enumerated list, the runtime should proceed rather than raise an error.

### Impact

`EnumeratedShapes` is documented as the recommended way to support multiple input shapes in a single CoreML model. It builds without error but fails at inference from any API. I found no workaround. The only viable path that I found is to compile separate fixed-shape models for each desired input size, which defeats the purpose of`EnumeratedShapes`.

This was discovered during conversion of [NX-AI/TiRex](https://huggingface.co/NX-AI/TiRex) (xLSTM-based time series forecasting model) to CoreML for on-device deployment on iPhone and Apple Watch.

### Current workaround

Compile separate fixed-shape models per configuration:

```python
for batch_size in [1, 4, 8, 16, 32]:
ct_input = ct.TensorType(
name="context",
shape=(batch_size, 256),
dtype=np.float32,
)
mlmodel = ct.convert(traced, inputs=[ct_input], ...)
mlmodel.save(f"model_b{batch_size}.mlpackage")
```

Start with the minimal Python conversion and inference reproduction using ct.EnumeratedShapes, then compare its batch-size results with the fixed-shape workaround. Exercise the listed Swift input-construction paths and inspect whether the failure is in conversion or the Apple runtime; done requires supported enumerated shapes to infer correctly without the reported stride errors.

Written by the indexing model from the issue text.