TopK fails to analyze for a small LSTM model

[alexander-camuto]

A smallish LSTM fails to unify DIM types (Impossible to unify TDim with I64) for a TOPK node:

Output from cli

┏ 0 Source input_ids
┃   ━━━ batch_size,64,I64
┣┻ 10 Gather /embedding/Gather
┃   ━━━ batch_size,64,32,F32
┃┏ 1 Source attention_mask
┃┃   ━━━ batch_size,64,F32
┃┣ 11 Reduce<Sum> /ReduceSum
┃┃   ━━━ batch_size,F32
┃┣ 12 onnx.Cast /Cast
┃┣ 13 onnx.Cast /Cast_1
┃┃   ━━━ batch_size,TDim
┃┣┓ 
┃┃┣ 15 Shape /Shape
┃┃┃   ━━━ 1,TDim batch_size
┃┃┣┻ 17 Gather /Gather
┃┣╋ 18 Topk /TopK
┃┃┃   ━━━ ?,?
┃┃┃   ━━━ ?,?
┃┃┣ 19 onnx.Cast /Cast_2
┃┃┣┓ 
┗━┓┃ 
┃┃┣┻ 20 Gather /Gather_1
┃┃┃   ━━━ ?,?,?,F32
┃┃┣ 21 PermuteAxes /Transpose
┗━┓ 
┃┃┣ 23 onnx.Cast /Cast_3
┃┣━┓ 
┃┃┃┣ 22 Shape /Shape_1
┃┃┃┃   ━━━ ?,TDim
┃┃┃┣┻ 34 Gather /lstm/Gather
┃┃┃┃   ━━━ TDim
┃┃┃┣ 36 AddDims /lstm/Unsqueeze
┃┃┃┃   ━━━ 1,TDim
┃┃┃┣┻┻ 38 InferenceConcat /lstm/Concat
┃┃┃┃   ━━━ 3,TDim
┃┃┃┣ 39 ConstantOfShape /lstm/ConstantOfShape
┃┃┃┃   ━━━ 2,?,32,F32
 
┃┃ 
┃┣╋╋┻┻┻┻ 40 LSTM /lstm/LSTM
┃┃    ━━━ ?,2,?,32,F32
┃┃    ━━━ 2,?,32,F32
┃┃    ━━━ 2,?,32,F32
┃┣ 41 PermuteAxes /lstm/Transpose
┃┃   ━━━ ?,?,?,?,F32
┃┣┻ 43 Reshape /lstm/Reshape
┃┃   ━━━ ?,?,?,F32
┃┣ 44 PermuteAxes /Transpose_1
┣━┓ 
┃┃┣ 30 Shape /Shape_3
┃┃┃   ━━━ ?,TDim
┃┃┣ 31 ConstantOfShape /ConstantOfShape
┃┃┃   ━━━ ..,I64
┣━━┓ 
┃┃┃┣ 24 Shape /Shape_2
┃┃┃┃   ━━━ ?,TDim
┃┃┃┣ 25 Reduce<Prod> /ReduceProd
┃┃┃┃   ━━━ TDim
┃┃┃┣ 27 onnx.Cast /Cast_4
┃┃┃┣┻┻ 29 Range /Range
┃┃┃┃   ━━━ ?,TDim
┗━┓┃ 
┃┣┻┻ 32 ScatterElements /ScatterElements
┃┃   ━━━ ?,?
┣┻ 45 Gather /Gather_2
┣┻ 46 Gather /Gather_3
┃   ━━━ ?,64,F32
┣┻┻ 47 Gemm /layers/layers.0/Gemm
┃   ━━━ ?,32,F32
┣ 48 Clip /layers/layers.1/Relu
┣┻┻ 49 Gemm /layers/layers.3/Gemm
    ━━━ ?,20,F32
[2025-02-06T18:37:57.164582000Z ERROR tract] Error at stage "analyse"
    
    Caused by:
        0: ModelBuildingError
        1: Failed analyse for node #18 "/TopK" Topk
        2: Infering facts
        3: Applying rule inputs[1].datum_type == I64
        4: Impossible to unify TDim with I64.

Model

network.onnx.zip

To reproduce

tract network.onnx