CoreML crashes with TensorFlow 2 models with multiple Conv1D and LayerNorm layers

main.swift

import CoreML

let model = try! MyModel()

let inputData = [Float](repeating: 0.1, count: 4096 * 2)
let inputArray = MLShapedArray(scalars: inputData, shape: [1, 4096, 2])
let input = MyModelInput(input_1: inputArray)

let output = try! model.prediction(input: input)
print("\(output.IdentityShapedArray.shape)")

make_model.py

# Python 3.9.12
# pip install tensorflow==2.6.2 coremltools==5.2.0 protobuf<=3.20

import coremltools as ct
import tensorflow as tf
import os


def make_model(conv_layer_definitions: list[int]) -> tf.keras.Model:
    input = tf.keras.layers.Input(shape=(4096, 2))
    output = input
    for conv_layer_definition in conv_layer_definitions:
        output = tf.keras.layers.Conv1D(conv_layer_definition, 8, 2, "same")(output)
        output = tf.keras.layers.LayerNormalization()(output)
    model = tf.keras.Model(inputs=input, outputs=output)
    model.compile(optimizer="SGD", loss="binary_crossentropy")
    model.summary()
    return model


conv_layer_definitions_not_working = [
    [32, 32, 64, 64, 128, 128, 256],
    [32, 32, 64, 64, 64, 64],
    [32, 32, 64, 64, 64],
    [128, 128, 128, 128],
    [256, 256, 256],
    [512, 512, 512],
    [4096],
]

conv_layer_definitions_works_well = [
    [32, 32, 64, 64],
    [64, 64, 64, 64],
    [128, 128, 128],
    [256, 256],
    [512, 512],
]
# replace this with conv_layer_definitions_not_working[...] and the Swift program below will crash as reported above.
conv_layer_definitions = conv_layer_definitions_works_well[-1]


os.system("rm -rf MyModel.mlpackage")

converted_model: ct.models.MLModel = ct.convert(
    make_model(conv_layer_definitions), convert_to="mlprogram", source="tensorflow"
)
converted_model.save("MyModel.mlpackage")

stacktrace

* thread #1, queue = 'com.apple.CoreMLBatchProcessingQueue', stop reason = EXC_BAD_ACCESS (code=1, address=0x0)
    frame #0: 0x00000001bb52711c libsystem_platform.dylib`_platform_memmove + 76
    frame #1: 0x00000001d0589b58 Espresso`EspressoLight::espresso_plan::__copy_inputs(std::__1::shared_ptr<EspressoLight::plan_task_t>, std::__1::shared_ptr<Espresso::abstract_batch> const&, int, std::__1::shared_ptr<Espresso::net>) + 1436
    frame #2: 0x00000001d0588e2c Espresso`EspressoLight::espresso_plan::dispatch_task_on_compute_batch(std::__1::shared_ptr<Espresso::abstract_batch> const&, std::__1::shared_ptr<EspressoLight::plan_task_t> const&) + 504
    frame #3: 0x00000001d0593a24 Espresso`EspressoLight::espresso_plan::execute_sync() + 412
    frame #4: 0x00000001d0599388 Espresso`espresso_plan_execute_sync + 132
    frame #5: 0x00000001c316bf40 CoreML`-[MLNeuralNetworkEngine executePlan:error:] + 136
    frame #6: 0x00000001c316c5e8 CoreML`-[MLNeuralNetworkEngine evaluateInputs:bufferIndex:options:error:] + 728
    frame #7: 0x00000001c316ead4 CoreML`__54-[MLNeuralNetworkEngine evaluateInputs:options:error:]_block_invoke + 44
    frame #8: 0x00000001005f63a8 libdispatch.dylib`_dispatch_client_callout + 20
    frame #9: 0x000000010060ab94 libdispatch.dylib`_dispatch_lane_barrier_sync_invoke_and_complete + 192
    frame #10: 0x00000001c316e8f8 CoreML`-[MLNeuralNetworkEngine evaluateInputs:options:error:] + 376
    frame #11: 0x00000001c3163568 CoreML`__62-[MLNeuralNetworkEngine predictionFromFeatures:options:error:]_block_invoke + 128
    frame #12: 0x00000001005f63a8 libdispatch.dylib`_dispatch_client_callout + 20
    frame #13: 0x000000010060ab94 libdispatch.dylib`_dispatch_lane_barrier_sync_invoke_and_complete + 192
    frame #14: 0x00000001c31633dc CoreML`-[MLNeuralNetworkEngine predictionFromFeatures:options:error:] + 436
  * frame #15: 0x000000010000d958 MyCoreMLCmdApp`MyModel.prediction(input=0x00000001010af390, options=0x00000001010ac080, self=0x00000001010cb250) at MyModel.swift:224:37
    frame #16: 0x000000010000d878 MyCoreMLCmdApp`MyModel.prediction(input=0x00000001010af390, self=0x00000001010cb250) at MyModel.swift:209:25
    frame #17: 0x00000001000045c8 MyCoreMLCmdApp`main at main.swift:43:25
    frame #18: 0x000000010004108c dyld`start + 520