CoreML gives unexpected output shape for a model with dynamic input shape

Hello.

I am manually constructing some models with the CoreML protobuf format. When the model has flexible input shapes, I am seeing unexpected output shapes in some cases after running prediction(from:).

The model is a single matrix multiplication, A*B (one innerProduct layer), and the dynamic dimension is the first dimension of the only input A (B is constant).

What I observe is that sometimes there are additional leading ones in the output shape.

Some test program output showing the shapes:

running model: dynamic_shape.mlmodel
A shape: [1, 2]
Y shape: [1, 1, 1, 1, 4]
running model: dynamic_shape.mlmodel
A shape: [2, 2]
Y shape: [1, 1, 1, 2, 4]
running model: dynamic_input_shape.mlmodel
A shape: [1, 2]
Y shape: [1, 4]
running model: dynamic_input_shape.mlmodel
A shape: [2, 2]
Y shape: [1, 1, 1, 2, 4]
running model: static_shape.mlmodel
A shape: [1, 2]
Y shape: [1, 4]

I've put the model generation and test code below. Am I specifying the dynamic input/output shapes correctly when creating the .mlmodel? Is the output shape given by CoreML expected, and if so, why are there leading ones? Would appreciate any input.

Python script to generate .mlmodel files. coremltools version is 6.3.0.

from coremltools.proto.Model_pb2 import Model
from coremltools.proto.FeatureTypes_pb2 import ArrayFeatureType
from coremltools.proto.NeuralNetwork_pb2 import EXACT_ARRAY_MAPPING

def build_model(with_dynamic_input_shape: bool, with_dynamic_output_shape: bool):
    model = Model()
    model.specificationVersion = 4

    input = model.description.input.add()
    input.name = "A"
    input.type.multiArrayType.shape[:] = [1, 2]
    input.type.multiArrayType.dataType = ArrayFeatureType.FLOAT32

    if with_dynamic_input_shape:
        range = input.type.multiArrayType.shapeRange.sizeRanges.add()
        range.upperBound = -1

        range = input.type.multiArrayType.shapeRange.sizeRanges.add()
        range.lowerBound = 2
        range.upperBound = 2

    output = model.description.output.add()
    output.name = "Y"
    output.type.multiArrayType.shape[:] = [1, 4]
    output.type.multiArrayType.dataType = ArrayFeatureType.FLOAT32
    
    if with_dynamic_output_shape:
        range = output.type.multiArrayType.shapeRange.sizeRanges.add()
        range.upperBound = -1

        range = output.type.multiArrayType.shapeRange.sizeRanges.add()
        range.lowerBound = 4
        range.upperBound = 4

    layer = model.neuralNetwork.layers.add()
    layer.name = "MatMul"
    layer.input[:] = ["A"]
    layer.output[:] = ["Y"]
    layer.innerProduct.inputChannels = 2
    layer.innerProduct.outputChannels = 4
    layer.innerProduct.weights.floatValue[:] = [0.0, 4.0, 1.0, 5.0, 2.0, 6.0, 3.0, 7.0]

    model.neuralNetwork.arrayInputShapeMapping = EXACT_ARRAY_MAPPING

    return model

if __name__ == "__main__":
    model = build_model(with_dynamic_input_shape=True, with_dynamic_output_shape=True)
    with open("dynamic_shape.mlmodel", mode="wb") as f:
        f.write(model.SerializeToString(deterministic=True))

    model = build_model(with_dynamic_input_shape=True, with_dynamic_output_shape=False)
    with open("dynamic_input_shape.mlmodel", mode="wb") as f:
        f.write(model.SerializeToString(deterministic=True))

    model = build_model(with_dynamic_input_shape=False, with_dynamic_output_shape=False)
    with open("static_shape.mlmodel", mode="wb") as f:
        f.write(model.SerializeToString(deterministic=True))

Swift program to run the models and print the output shape.

import Foundation
import CoreML

func makeFloatShapedArray(shape: [Int]) -> MLShapedArray<Float> {
    let size = shape.reduce(1, *)
    let values = (0 ..< size).map { Float($0) }
    return MLShapedArray(scalars: values, shape: shape)
}

func runModel(model_path: URL, m: Int) throws {
    print("running model: \(model_path.lastPathComponent)")
    let compiled_model_path = try MLModel.compileModel(at: model_path)
    let model = try MLModel(contentsOf: compiled_model_path)
    
    let a = MLMultiArray(makeFloatShapedArray(shape: [m, 2]))
    print("A shape: \(a.shape)")
    let inputs = try MLDictionaryFeatureProvider(dictionary: ["A": a])
    
    let outputs = try model.prediction(from: inputs)
    let y = outputs.featureValue(for: "Y")!.multiArrayValue!
    
    print("Y shape: \(y.shape)")
}

func modelUrl(_ model_file: String) -> URL {
    return URL(filePath: "/path/to/models/\(model_file)")
}

try runModel(model_path: modelUrl("dynamic_shape.mlmodel"), m: 1)
try runModel(model_path: modelUrl("dynamic_shape.mlmodel"), m: 2)
try runModel(model_path: modelUrl("dynamic_input_shape.mlmodel"), m: 1)
try runModel(model_path: modelUrl("dynamic_input_shape.mlmodel"), m: 2)
try runModel(model_path: modelUrl("static_shape.mlmodel"), m: 1)