Handle Model Display Correctly

[heywhy]

Fixes #604 and multi inputs model, see example below:

[Axon.input("input_1"), Axon.input("input_2")]
|> Axon.concatenate()
|> Axon.dense(32, activation: :relu)
|> Axon.dropout(rate: 0.5)
|> Axon.dense(1)
|> Axon.Display.as_table(%{"input_1" => Nx.iota({2, 1}), "input_2" => Nx.iota({2, 2})})

Expected output

+------------------------------------------------------------------------------------------------------------------------------------------+
|                                                                  Model                                                                   |
+===============================+============================================+======================+=================+====================+
| Layer                         | Input Shape                                | Output Shape         | Options         | Parameters         |
+===============================+============================================+======================+=================+====================+
| input_1 ( input )             | %{}                                        | s32[2][1]            | shape: nil      |                    |
|                               |                                            |                      | optional: false |                    |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+
| input_2 ( input )             | %{}                                        | s32[2][2]            | shape: nil      |                    |
|                               |                                            |                      | optional: false |                    |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+
| container_0 ( container )     | {}                                         | s32[2][1], s32[2][2] |                 |                    |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+
| concatenate_0 ( concatenate ) | %{"container_0" => "s32[2][1], s32[2][2]"} | s32[2][3]            | axis: -1        |                    |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+
| dense_0 ( dense )             | %{"concatenate_0" => "s32[2][3]"}          | f32[2][32]           |                 | kernel: f32[3][32] |
|                               |                                            |                      |                 | bias: f32[32]      |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+
| relu_0 ( relu )               | %{"dense_0" => "f32[2][32]"}               | f32[2][32]           |                 |                    |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+
| dropout_0 ( dropout )         | %{"relu_0" => "f32[2][32]"}                | f32[2][32]           | rate: 0.5       | key: f32[2]        |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+
| dense_1 ( dense )             | %{"dropout_0" => "f32[2][32]"}             | f32[2][1]            |                 | kernel: f32[32][1] |
|                               |                                            |                      |                 | bias: f32[1]       |
+-------------------------------+--------------------------------------------+----------------------+-----------------+--------------------+