`Bias` fails to broadcast in the context of `matmul` in tf lite model

YaoJiayi commented 1 year ago

1. System information

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04
TensorFlow installation (pip package or built from source): pip
TensorFlow library (version, if pip package or github SHA, if built from source): 2.14.0-dev20230602

2. Code

This is the minimized code to reproduce the issue:

import tensorflow as tf
import numpy as np

x1 = tf.constant([1., 2.], shape=[1, 2])

class Model(tf.keras.Model):
  def __init__(self):
    super(Model, self).__init__()
    self.w = tf.Variable([[3., 4.], [5., 6.]])
    self.b = tf.Variable([3.])
  @tf.function(input_signature=[tf.TensorSpec(x1.shape, x1.dtype)])
  def call(self, x):
    return tf.matmul(x, self.w) + self.b

m = Model()
print('Keras mode output: ', m(x1).numpy())

converter = tf.lite.TFLiteConverter.from_keras_model(m)
tflite_model = converter.convert()
def _evaluateTFLiteModel(tflite_model, input_data):
    interpreter = tf.lite.Interpreter(model_content=tflite_model)
    interpreter.allocate_tensors()

    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    for i in range(len(input_data)):
        interpreter.set_tensor(input_details[i]['index'], input_data[i])

    interpreter.invoke()

    output_data = [interpreter.get_tensor(output_details[i]['index'])
                   for i in range(len(output_details))]
    return output_data

print('Lite mode output: ', _evaluateTFLiteModel(tflite_model,[x1])[0])

3. Failure after conversion

Output:

Keras mode output: [[16. 19.]]
Lite mode output:  RuntimeError: tensorflow/lite/kernels/fully_connected.cc:360 NumElements(bias) != SizeOfDimension(filter, 0) (1 != 2)Node number 0 (FULLY_CONNECTED) failed to prepare.Failed to apply the default TensorFlow Lite delegate indexed at 0.

Conversion Failure:

During conversion, the model fails due to the following check at tensorflow/lite/kernels/fully_connected.cc:360
```
if (bias) {
TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(filter, 0));
}
```