Removing specific versions or updating to the latest available version corrected the issue for me. Pol Valente was kind enough to update the Axon lib, thanks Pol!
Mix.install([
:axon,
:nx,
:exla,
:kino,
:table_rex,
:scidata
])
Nx.default_backend(EXLA.Backend)
…
…
…
Axon.Display.as_table(model, template)
|> IO.puts()
+------------------------------------------------------------------------------------------------------------------+
| Model |
+=======================+===============================+==============+===================+=======================+
| Layer | Input Shape | Output Shape | Options | Parameters |
+=======================+===============================+==============+===================+=======================+
| images ( input ) | %{} | f32[1][784] | shape: {nil, 784} | |
| | | | optional: false | |
+-----------------------+-------------------------------+--------------+-------------------+-----------------------+
| dense_0 ( dense ) | %{"images" => "f32[1][784]"} | f32[1][128] | | kernel: f32[784][128] |
| | | | | bias: f32[128] |
+-----------------------+-------------------------------+--------------+-------------------+-----------------------+
| relu_0 ( relu ) | %{"dense_0" => "f32[1][128]"} | f32[1][128] | | |
+-----------------------+-------------------------------+--------------+-------------------+-----------------------+
| dense_1 ( dense ) | %{"relu_0" => "f32[1][128]"} | f32[1][10] | | kernel: f32[128][10] |
| | | | | bias: f32[10] |
+-----------------------+-------------------------------+--------------+-------------------+-----------------------+
| softmax_0 ( softmax ) | %{"dense_1" => "f32[1][10]"} | f32[1][10] | | |
+-----------------------+-------------------------------+--------------+-------------------+-----------------------+
Total Parameters: 101770
Total Parameters Memory: 407.08 kilobytes