keras model with tile + concatenate + Conv2D gives "StatefulPartitionedCall is not supported"

[nikoladze]

I got "StatefulPartitionedCall is not supported" when trying to convert a model created with keras and Conv2D layers inside. The onnx model is created, but can't be run because of the StatefulPartitionedCall op somewhere inside. I was able to isolate the issue to a particular chain of tile + concatenate + Conv2D

The particular use case where this showed up can work around this by replacing the Conv2D layers by Dense layers since it were 1x1 convolutions anyways. For this combination the issue does not appear.

System information

• Arch Linux
• TensorFlow Version: 2.19.1
• Python version: 3.12
• ONNX version (if applicable, e.g. 1.11*): 1.17.0
• ONNXRuntime version (if applicable, e.g. 1.11*): 1.22.2
• tf2onnx from https://github.com/onnx/tensorflow-onnx#4fed7de9534b6a084f7f2326bae775545bd97f9e

To Reproduce

from keras import ops
from keras.layers import Conv2D, Input
from keras.models import Model
import tensorflow as tf

def get_model():
    inp1 = Input((32, 1, 3))
    inp2 = Input((32, 32, 3))
    x1 = ops.tile(inp1, (1, 1, 32, 1))
    x2 = inp2
    x = ops.concatenate([x1, x2], axis=-1)
    x = Conv2D(32, 1)(x)
    return Model([inp1, inp2], x)

model = get_model()
model([tf.random.uniform((1, 32, 1, 3)), tf.random.uniform((1, 32, 32, 3))])
model.export("model.onnx", format="onnx")

gives

...
Tensorflow op [functional_2_1[/conv2d_3_1/StatefulPartitionedCall](http://localhost:8889/conv2d_3_1/StatefulPartitionedCall): StatefulPartitionedCall] is not supported
Unsupported ops: Counter({'StatefulPartitionedCall': 1})
...

This was the most minimal example i found. There is no issue without the tile operation, e.g. already feeding inp1 in the same shape as inp2 and just doing concatenate.

Exporting to SavedModel format first and running tf2onnx on the command line gives the same error. I tried different opset versions and --opset 18 was the highest available, but gives the same error.