tensorflow: TFLite Concatenation Fails on GPU delegate

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes.
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Mac
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: Tested on Samsung Galaxy
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): v2.1.0-rc2-17-ge5bf8de410 2.1.0
• Python version: 3.7
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:
• GPU model and memory:

Describe the current behavior Concatenation Op for non BCHW dimensions fails on GPU delegate, same model runs fine on CPU delegate.

(See below code for creating model.tflite)

The following benchmark runs successfully (CPU):

adb shell /data/local/tmp/benchmark_model      --graph=/data/local/tmp/model.tflite     --num_threads=4 --use_gpu=false --num_runs=1000

The following benchmark run fails (GPU):

adb shell /data/local/tmp/benchmark_model      --graph=/data/local/tmp/model.tflite     --num_threads=4 --use_gpu=true --num_runs=1000

Error message:

INFO: Initialized TensorFlow Lite runtime.
INFO: Created TensorFlow Lite delegate for GPU.
ERROR: TfLiteGpuDelegate Init: CONCATENATION: Dimensions are not BHWC: 2 1 30
ERROR: TfLiteGpuDelegate Prepare: delegate is not initialized
ERROR: Node number 1 (TfLiteGpuDelegateV2) failed to prepare.

Describe the expected behavior

TFLite Model should run concat op succesfully both on CPU and GPU.

Standalone code to reproduce the issue

import numpy as np
import tensorflow as tf

# Construct a basic model.
root = tf.train.Checkpoint()
root.v1 = tf.Variable(np.ones((1, 10), dtype=np.float32))
root.v2 = tf.Variable(np.ones((1, 20), dtype=np.float32))
root.f = tf.function(lambda x: tf.concat([root.v1, root.v2, x], axis=1))

# Save the model.
export_dir = "/tmp/test_saved_model"
input_data = tf.constant(np.ones((1, 30), dtype=np.float32))
to_save = root.f.get_concrete_function(input_data)
tf.saved_model.save(root, export_dir, to_save)

# Convert the model.
converter = tf.lite.TFLiteConverter.from_saved_model(export_dir)
tflite_model = converter.convert()
open("model.tflite", "wb").write(tflite_model)