fix: fix Axon.Display for Tensor parameter templates

lib/axon/display.ex

@@ -186,6 +186,9 @@ defmodule Axon.Display do
         %Parameter{shape: {:tuple, shapes}}, acc ->
           Enum.reduce(shapes, acc, &(Nx.size(apply(&1, input_shapes)) + &2))
 
+        %Parameter{template: %Nx.Tensor{} = template}, acc ->
+          acc + Nx.size(template)
+
         %Parameter{template: shape_fn}, acc when is_function(shape_fn) ->
           acc + Nx.size(apply(shape_fn, input_shapes))
       end)
@@ -253,6 +256,11 @@ defmodule Axon.Display do
 
         "#{name}: tuple#{inspect(shapes)}"
 
+      %Parameter{name: name, template: %Nx.Tensor{} = template} ->
+        type = Nx.type(template)
+        shape = Nx.shape(template)
+        "#{name}: #{type_str(type)}#{shape_string(shape)}"
+
       %Parameter{name: name, template: shape_fn} when is_function(shape_fn) ->
         shape = Nx.shape(apply(shape_fn, input_shapes))
         "#{name}: #{type_str(type)}#{shape_string(shape)}"

test/axon/display_test.exs

@@ -0,0 +1,18 @@
+defmodule Axon.DisplayTest do
+  use ExUnit.Case, async: true
+
+  describe "as_table/2" do
+    test "renders layers with concrete parameter templates" do
+      model = Axon.input("input") |> Axon.dense(32)
+      input = Nx.template({1, 16}, :f32)
+
+      table = Axon.Display.as_table(model, input)
+
+      assert table =~ ~s|dense_0 ( dense )|
+      assert table =~ ~s|kernel: f32[16][32]|
+      assert table =~ ~s|bias: f32[32]|
+      assert table =~ ~s|Total Parameters: 544|
+      assert table =~ ~s|Total Parameters Memory: 2.18 kilobytes|
+    end
+  end
+end