Issue with GatherV2 conversion

[pziecina]

Describe the bug
After converting toy model with below structure:

input_a = tf.keras.layers.Input(shape=(length,), name="input_a", dtype="float32")
input_b = tf.keras.layers.Input(shape=(length, length * 2), name="input_b", dtype="float32")

input_a_log_probs = tf.nn.log_softmax(input_a, axis=-1, name="input_a_log_probs")
input_a_top_log_probs, input_a_top_index = tf.math.top_k(input_a_log_probs, k=k, name="input_a_top_index")
gathered = tf.gather(input_b, input_a_top_index, axis=1, batch_dims=1)
expanded = tf.expand_dims(gathered, axis=1)
output = tf.broadcast_to(expanded, [batch_size, length, k, length * 2])

inputs = [input_a, input_b]
outputs = [output]
model = tf.keras.Model(inputs=inputs, outputs=outputs)

I'm trying to infer it on onnxruntime but receiving below error:

2021-02-04 14:57:34.647032243 [I:onnxruntime:, inference_session.cc:223 operator()] Flush-to-zero and denormal-as-zero are off
2021-02-04 14:57:34.647134789 [I:onnxruntime:, inference_session.cc:230 ConstructorCommon] Creating and using per session threadpools since use_per_session_threads_ is true
Traceback (most recent call last):
  File "run_onnx.py", line 43, in <module>
    main()
  File "run_onnx.py", line 36, in main
    session = onnxruntime.InferenceSession(model_payload, providers=None, sess_options=session_options)
  File "/usr/local/lib/python3.8/dist-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 206, in __init__
    self._create_inference_session(providers, provider_options)
  File "/usr/local/lib/python3.8/dist-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 228, in _create_inference_session
    sess = C.InferenceSession(session_options, self._model_bytes, False, self._read_config_from_model)
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Node:PartitionedCall/functional_1/tf_op_layer_BroadcastTo/BroadcastTo Output:Identity:0 [ShapeInferenceError] Mismatch between number of source and target dimensions. Source=5 Target=4

I think that there might be wrongly converted GatherV2 op with batch_dims because when dumped its shape there appeared additional axis equal to batch size - (8, 8, 5, 758). On TF same tensor has shape (8, 5, 758)

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• Tensorflow Version: 2.3.1
• Python version: 3.8.5

To Reproduce

• Download broadcast_to_issue_bug_report.tar.gz and extract it in $PWD
• Start TF2 container with below command

docker run --rm -it -v $PWD:$PWD -w $PWD nvcr.io/nvidia/tensorflow:20.12-tf2-py3

• Init env, dump tf model, perform conversion and run inference running below script in container

pip install onnxruntime==1.6.0
pip install tf2onnx==1.8.2
python prepare_tf_model.py 2>&1|tee prepare_tf_model.log
python -m tf2onnx.convert --saved-model model.savedmodel --output model.onnx --opset 13 --verbose 2>&1|tee tf2onnx.log
python run_tf.py 2>&1|tee run_tf.log
python run_onnx.py 2>&1|tee run_onnx.log

[TomWildenhain-Microsoft]

Hi @pawel-ziecina, thanks for very thorough info to reproduce this bug. The issue here is because our GatherV2 conversion does not support the batch_dim property. I can relatively easily implement a conversion for batch dim in opset 12. If you need a lower opset, let me know what the highest acceptable opset is and I can implement it, though it will be a bit harder. We should definitely handle this case better and raise an error at the least.

[pziecina]

For my uses, fix for opset 12+ will do the job. Thanks!