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

[Suad0]

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<Float> + 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)

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)

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)

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<Float>(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 ofEnumeratedShapes.

This was discovered during conversion of 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:

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")

[metehankaygsz]

I tested the reproduction with coremltools 9.0 on macOS 26.5.1 (build 25F80), including PyTorch 2.11.0.

After saving and reloading the model, batch sizes 1, 4, 8, 16, and 32 all produced the expected output shapes and values. I also tested native Swift inference with MLMultiArray and MLShapedArray, plus torch.stack layouts similar to TiRex’s recurrent state and output paths. Those passed as well.

The reported Stack_nd error cannot originate from the posted x * 2 model because that graph contains no stack operation. This suggests the failing model differs from the minimal reproduction. The E5RT/Espresso stride error also comes from the operating-system Core ML runtime rather than the converter source in this repository.

Could you retest on macOS 26.5.1? If it still fails, please attach:

• The exact failing .mlpackage
• A minimal Swift project that reproduces the error
• The Python conversion script used to generate that package
• The output of sw_vers and the precise device OS version

With those artifacts, we can determine whether the issue is specific to the generated TiRex graph, an older runtime build, or a converter-generated model representation. I have prepared local regression coverage for the posted leading-dimension case, but there is not currently enough evidence for a source-level fix.