gather_nd fails when using varaiable batch dims

[alexander-camuto]

Can generate gather_nd ops from tensorflow as such:

import tf2onnx
import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model

# gather_nd in tf then export to onnx
x = in1 = Input((15, 18,))
w = in2 = Input((15, 1), dtype=tf.int32)
x = tf.gather_nd(x, w, batch_dims=1)
tm = Model((in1, in2), x )
tm.summary()
tm.compile(optimizer='adam', loss='mse')

spec = tf.TensorSpec([None, 15, 18], tf.float32, name='input')
index_spec = tf.TensorSpec( [None, 15, 1], tf.int32, name='index')

model_path = "network.onnx"

tf2onnx.convert.from_keras(tm, input_signature=[spec, index_spec], inputs_as_nchw=['input', 'index'], opset=13, output_path=model_path)

The resulting graph has variable batch sizes for the inputs:

Tract fails to analyze with:

┏ 0 Source input
┃   ━━━ unk__58,15,18,F32
┣┓ 
┃┣ 9 Shape model/tf.compat.v1.gather_nd/BatchGatherND/Shape
┃┃   ━━━ 3,TDim unk__58, 15, 18
┃┣ 10 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Shape__7
┃┃   ━━━ 3,I32
┃┣┻┻┻ 13 StridedSlice model/tf.compat.v1.gather_nd/BatchGatherND/strided_slice
┃┃   ━━━ 1,I32
┃┣┓ 
┃┃┣ 14 Split model/tf.compat.v1.gather_nd/BatchGatherND/unstack
┃┃┣┻ 15 Squeeze13 model/tf.compat.v1.gather_nd/BatchGatherND/unstack__11
┃┃┃   ━━━ I32
┃┃┣┻┻ 16 Range model/tf.compat.v1.gather_nd/BatchGatherND/Range
┃┃┃   ━━━ ?,I32
┃┃┣ 17 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Cast
┃┃┣┓  
┃┃┃┣┻ 28 Reshape model/tf.compat.v1.gather_nd/BatchGatherND/meshgrid/Reshape
┃┃┃┃   ━━━ ..,I32
┃┃┗┓ 
┃┃┃┣ 18 Size model/tf.compat.v1.gather_nd/BatchGatherND/meshgrid/Size
┃┃┃┃   ━━━ TDim
┃┃┃┣ 19 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/meshgrid/Size__13
┃┃┃┃   ━━━ I32
┃┃┃┣┻ 21 Unsqueeze13 model/tf.compat.v1.gather_nd/BatchGatherND/meshgrid/ones/packed
┃┃┃┃   ━━━ 1,I32
┃┃┃┣ 22 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/meshgrid/ones__25
┃┃┃┃   ━━━ 1,TDim
┃┃┃┣┻ 23 MultiBroadcastTo model/tf.compat.v1.gather_nd/BatchGatherND/meshgrid/ones
┃┃┃┃   ━━━ ..,I32
┃┃┣┻ 29 Mul model/tf.compat.v1.gather_nd/BatchGatherND/meshgrid/mul
┃┃┣┻ 30 Reshape model/tf.compat.v1.gather_nd/BatchGatherND/Reshape
┃┃┣┻ 31 Unsqueeze13 model/tf.compat.v1.gather_nd/BatchGatherND/stack
┃┃┃   ━━━ ?,?,I32
┃┃┣ 32 PermuteAxes model/tf.compat.v1.gather_nd/BatchGatherND/transpose
┃┃┣┓ 
┃┃┃┣ 33 Shape model/tf.compat.v1.gather_nd/BatchGatherND/Shape_2
┃┃┃┃   ━━━ ?,TDim
┃┃┃┣ 34 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Shape_2__30
┃┃┃┃   ━━━ ?,I32
┃┃┃┣┓    
┃┃┃┃┣┻┻┻ 35 StridedSlice model/tf.compat.v1.gather_nd/BatchGatherND/strided_slice_2
┃┃┃┗┓    
┃┃┃┃┣┻┻┻ 40 StridedSlice model/tf.compat.v1.gather_nd/BatchGatherND/strided_slice_3
 
 ┃ 
┃┃┃┣┻┻ 41 InferenceConcat model/tf.compat.v1.gather_nd/BatchGatherND/concat
┃┃┃┣ 42 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Reshape_1__37
┃┃┃┃   ━━━ ?,TDim
┃ 
┃┃┣┻ 43 Reshape model/tf.compat.v1.gather_nd/BatchGatherND/Reshape_1
┃┃┃   ━━━ ?,?,?,I32
┃┃┃┏ 1 Source index
┃┃┃┃   ━━━ unk__59,15,1,I32
┃┃┃┣┓ 
┃┃┃┃┣ 11 Shape model/tf.compat.v1.gather_nd/BatchGatherND/Shape_1
┃┃┃┃┃   ━━━ 3,TDim unk__59, 15, 1
┃┃┃┃┣ 12 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Shape_1__15
┃┃┃┃┃   ━━━ 3,I32
┃┃┃┃┣┓    
┃┃┃┃┃┣┻┻┻ 25 StridedSlice model/tf.compat.v1.gather_nd/BatchGatherND/strided_slice_1
┃┃┃┃┃┃   ━━━ 1,I32
┃┃┃┃┃┣┻┻ 26 InferenceConcat model/tf.compat.v1.gather_nd/BatchGatherND/concat_1
┃┃┃┃┃┃   ━━━ 3,I32
┃┃┃┃┃┣ 27 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Tile__38
┃┃┃┃┃┃   ━━━ 3,TDim
┃┃┗━┓┃ 
┃┃┃┃┣┻ 44 Tile model/tf.compat.v1.gather_nd/BatchGatherND/Tile
┃┣━━━┓ 
┃┃┃┃┃┣ 20 Reduce<Prod> model/tf.compat.v1.gather_nd/BatchGatherND/Prod
┃┃┃┃┃┃   ━━━ I32
┃┃┃┃┃┣┻ 24 Unsqueeze13 model/tf.compat.v1.gather_nd/BatchGatherND/concat_2/values_0
┃┃┃┃┃┃   ━━━ 1,I32
┃┃┃┗━┓    
┃┃┃┃┃┣┻┻┻ 36 StridedSlice model/tf.compat.v1.gather_nd/BatchGatherND/strided_slice_4
┃┃┃┃┃┃   ━━━ 2,I32
┃┃┃┃┣┻ 37 InferenceConcat model/tf.compat.v1.gather_nd/BatchGatherND/concat_2
┃┃┃┃┃   ━━━ 3,I32
┃┃┃┃┣ 38 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Reshape_2__42
┃┃┃┃┃   ━━━ 3,TDim
┃┃┗┓┃ 
┃┃┃┣┻ 39 Reshape model/tf.compat.v1.gather_nd/BatchGatherND/Reshape_2
┃┃┃┃   ━━━ ?,?,?,I32
┃┃┣┻ 45 InferenceConcat model/tf.compat.v1.gather_nd/BatchGatherND/concat_3
┃┃┣ 46 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/GatherNd__43
┃┃┃   ━━━ ?,?,?,TDim
┗┓┃ 
┃┣┻ 47 GatherNd model/tf.compat.v1.gather_nd/BatchGatherND/GatherNd
┃┃   ━━━ ?,?,..,F32
┃┣┓ 
┃┃┣ 48 Shape model/tf.compat.v1.gather_nd/BatchGatherND/Shape_3
┃┃┃   ━━━ ?,TDim
┃┃┣ 49 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Shape_3__45
┃┃┃   ━━━ ?,I32
┃┃┣┻┻┻ 50 StridedSlice model/tf.compat.v1.gather_nd/BatchGatherND/strided_slice_5
┗┓┃ 
┃┣┻ 51 InferenceConcat model/tf.compat.v1.gather_nd/BatchGatherND/concat_4
┃┃   ━━━ ?,I32
┃┣ 52 onnx.Cast model/tf.compat.v1.gather_nd/BatchGatherND/Reshape_3__49
┃┃   ━━━ ?,TDim
┣┻ 53 Reshape model/tf.compat.v1.gather_nd/BatchGatherND/Reshape_3
    ━━━ ..,F32
[2024-03-11T11:50:57.595587000Z ERROR tract] Error at stage "analyse"
    
    Caused by:
        0: ModelBuildingError
        1: #16 "model/tf.compat.v1.gather_nd/BatchGatherND/Range" Range has incomplete typing

Zipped file:

network.onnx.zip