Ordering of inputs to Scan nodes

[alexander-camuto]

I currently have an RNN of the form:

┏ 0 Source input
┃   ━━━ batch_size,3,F32
┣┻ 2 EinSum /RNN.extracted./RNN.Xt_WiT
┃   ━━━ batch_size,1,3,F32
┣╋ 4 Scan /RNN
   [loop] ┏ 0 Source /RNN.extracted./RNN.Xt_WiT
   [loop] ┃   ━━━ 1,1,3,F32
   [loop] ┣ 1 RmAxis /RNN.Xt_WiT.prop_axis.a.output
   [loop] ┃   ━━━ 1,3,F32
   [loop] ┃┏ 2 Source h_source
   [loop] ┃┃   ━━━ 1,3,F32
   [loop] ┃┣┻ 4 EinSum /RNN.Ht_1_RiT
   [loop] ┣┻ 5 Add /RNN.ht0
   [loop] ┣┻ 7 Add /RNN.ht_bias
   [loop] ┣ 8 Tanh /RNN.Ht
   [loop]     ━━━ 1,3,F32
     ━━━ batch_size,3,F32
     ━━━ 1,3,F32

If we inspect the Scan node more closely it is as such:

Node { id: 4, name: "/RNN", inputs: [
        3/0>,
        2/0>
    ], op: Scan { skip: 0, reset_every_turn: false, body: Graph { nodes: [Node { id: 0, name: "/RNN.extracted./RNN.Xt_WiT", inputs: [], op: TypedSource { fact: 1,
                        1,
                        3,F32
                    }, outputs: [
                        1,
                        1,
                        3,F32 >1/0
                    ]
                }, Node { id: 1, name: "/RNN.Xt_WiT.prop_axis.a.output", inputs: [
                        0/0>
                    ], op: Rm(0), outputs: [
                        1,
                        3,F32 >5/0
                    ]
                }, Node { id: 2, name: "h_source", inputs: [], op: TypedSource { fact: 1,
                        3,F32
                    }, outputs: [
                        1,
                        3,F32 >4/0
                    ]
                }, Node { id: 3, name: "R", inputs: [], op: Const(3,
                    3,F32 -0.5741769,
                    0.19402182,
                    -0.52662057,
                    -0.3177815,
                    0.48619103,
                    0.41249704,
                    -0.15882403,
                    0.34020323,
                    0.118276834), outputs: [
                        3,
                        3,F32 -0.5741769,
                        0.19402182,
                        -0.52662057,
                        -0.3177815,
                        0.48619103,
                        0.41249704,
                        -0.15882403,
                        0.34020323,
                        0.118276834 >4/1
                    ]
                }, Node { id: 4, name: "/RNN.Ht_1_RiT", inputs: [
                        2/0>,
                        3/0>
                    ], op: EinSum ak,nk->bn (F32), outputs: [
                        1,
                        3,F32 >5/1
                    ]
                }, Node { id: 5, name: "/RNN.ht0", inputs: [
                        1/0>,
                        4/0>
                    ], op: TypedBinOp(Add), outputs: [
                        1,
                        3,F32 >7/0
                    ]
                }, Node { id: 6, name: "/RNN.bi", inputs: [], op: Const(1,
                    3,F32 -0.33524436,
                    0.5043155,
                    0.41468257), outputs: [
                        1,
                        3,F32 -0.33524436,
                        0.5043155,
                        0.41468257 >7/1
                    ]
                }, Node { id: 7, name: "/RNN.ht_bias", inputs: [
                        5/0>,
                        6/0>
                    ], op: TypedBinOp(Add), outputs: [
                        1,
                        3,F32 >8/0
                    ]
                }, Node { id: 8, name: "/RNN.Ht", inputs: [
                        7/0>
                    ], op: ElementWiseOp(Tanh), outputs: [
                        1,
                        3,F32
                    ]
                }
            ], inputs: [
                2/0>,
                0/0>
            ], outputs: [
                8/0>
            ], outlet_labels: {}, properties: {}, symbol_table:
        }, decluttered: true, input_mapping: [State, Scan(ScanInfo { axis: 0, chunk: 1
            })
        ], output_mapping: [State. Last value to outlet 1. Full value to outlet 0 (axis: 0).
        ]
    }, outputs: [
        1,
        3,F32 ,
        1,
        3,F32
    ]
}

The two inputs to the scan node are 3/0> and 2/0.

Outlet 3/0 (not seen) is a Const node of shape [1,3]. Whereas (2/0) is of size [batch_size,1,3]. And the input_mapping makes sense to me in that it has order [State, Scan(ScanInfo { axis: 0, chunk: 1})]: 3/0 is the initial state node, and (2/0) is the input we slice over at each iteration of the Scan loop.

However once we look at the Scan node's subgraph the sources are in the following order:

 [loop] ┏ 0 Source /RNN.extracted./RNN.Xt_WiT
 [loop] ┃   ━━━ 1,1,3,F32
...
 [loop] ┃┏ 2 Source h_source
 [loop] ┃┃   ━━━ 1,3,F32

It sort of seems like:

• (Claim A) Source 0 should be matched with outlet (2/0) (of size [batch_size,1,3] which when sliced is of size [1,1,3]).
• (Claim B) Source 2 should be matched with outlet (3/0) (of size [1,3] and State).

However the ordering of sources and inputs here seems to imply otherwise.

I was wondering if Claim A and B are correct here and if so where / how are the wrapping node's inputs matched to the corresponding Scan Subgraph's sources ?

Attached you'll find the RNN in question:

network.onnx.zip

[kali]

You can get more out of the command line with a couple of flags:

NUC 30/08 11:01 ~/dev/sonos/tract/issue-1173% cargo run -p tract --  network.onnx dump --io-long --info --const
    Finished dev [unoptimized + debuginfo] target(s) in 0.13s
     Running `/home/kali/dev/sonos/tract/target/debug/tract network.onnx dump --io-long --info --const`
┏ 0 Source input
┃   * output fact #0: batch_size,3,F32 >2/0 MODEL INPUT #0
┃┏ 1 Const /RNN.extracted./RNN.Xt_WiT.prop_axis.a.input_1
┃┃   * output fact #0: 3,3,F32 -0.37928146, 0.113702476, 0.5145674, -0.11357668, 0.08259404, -0.47074276, -0.09194785, -0.28733814, 0.1228261 >2/1
┣┻ 2 EinSum /RNN.extracted./RNN.Xt_WiT
┃   * input fact  #0: 0/0> batch_size,3,F32
┃     input fact  #1: 1/0> 3,3,F32 -0.37928146, 0.113702476, 0.5145674, -0.11357668, 0.08259404, -0.47074276, -0.09194785, -0.28733814, 0.1228261
┃   * output fact #0: batch_size,1,3,F32 >4/1
┃   * ak,nk->abn (F32)
┃┏ 3 Const /RNN.h_chunk
┃┃   * output fact #0: 1,3,F32 0, 0, 0 >4/0
┗┓
┣╋ 4 Scan /RNN
     * input fact  #0: 3/0> 1,3,F32 0, 0, 0
       input fact  #1: 2/0> batch_size,1,3,F32
     * output fact #0: batch_size,3,F32  MODEL OUTPUT #0 output
       output fact #1: 1,3,F32  MODEL OUTPUT #1 25
     * Model input  #0: State
     * Model input  #1: Scan(ScanInfo { axis: 0, chunk: 1 })
     * Model output #0: State. Last value to outlet 1. Full value to outlet 0 (axis: 0).
     * skip:0 reset_every_turn:false
   [loop] ┏ 0 Source /RNN.extracted./RNN.Xt_WiT
   [loop] ┃   * output fact #0: 1,1,3,F32 >1/0 MODEL INPUT #1
   [loop] ┣ 1 RmAxis /RNN.Xt_WiT.prop_axis.a.output
   [loop] ┃   * input fact  #0: 0/0> 1,1,3,F32
   [loop] ┃   * output fact #0: 1,3,F32 >5/0
   [loop] ┃   * Axis: 0
   [loop] ┃┏ 2 Source h_source
   [loop] ┃┃   * output fact #0: 1,3,F32 >4/0 MODEL INPUT #0
   [loop] ┃┃┏ 3 Const R
   [loop] ┃┃┃   * output fact #0: 3,3,F32 -0.5741769, 0.19402182, -0.52662057, -0.3177815, 0.48619103, 0.41249704, -0.15882403, 0.34020323, 0.118276834 >4/1
   [loop] ┃┣┻ 4 EinSum /RNN.Ht_1_RiT
   [loop] ┃┃   * input fact  #0: 2/0> 1,3,F32
   [loop] ┃┃     input fact  #1: 3/0> 3,3,F32 -0.5741769, 0.19402182, -0.52662057, -0.3177815, 0.48619103, 0.41249704, -0.15882403, 0.34020323, 0.118276834
   [loop] ┃┃   * output fact #0: 1,3,F32 >5/1
   [loop] ┃┃   * ak,nk->bn (F32)
   [loop] ┣┻ 5 Add /RNN.ht0
   [loop] ┃   * input fact  #0: 1/0> 1,3,F32
   [loop] ┃     input fact  #1: 4/0> 1,3,F32
   [loop] ┃   * output fact #0: 1,3,F32 >7/0
   [loop] ┃┏ 6 Const /RNN.bi
   [loop] ┃┃   * output fact #0: 1,3,F32 -0.33524436, 0.5043155, 0.41468257 >7/1
   [loop] ┣┻ 7 Add /RNN.ht_bias
   [loop] ┃   * input fact  #0: 5/0> 1,3,F32
   [loop] ┃     input fact  #1: 6/0> 1,3,F32 -0.33524436, 0.5043155, 0.41468257
   [loop] ┃   * output fact #0: 1,3,F32 >8/0
   [loop] ┣ 8 Tanh /RNN.Ht
   [loop]     * input fact  #0: 7/0> 1,3,F32
   [loop]     * output fact #0: 1,3,F32  MODEL OUTPUT #0

Source 0 is actually model input 1, while Source 2 is model input 0. The mapping is match in model input order, not the node ids.

Hope it helps

[alexander-camuto]

gotcha ! thank you :)