diff --git a/lib/axon/display.ex b/lib/axon/display.ex index 057ce4fb..8bfdf4e7 100644 --- a/lib/axon/display.ex +++ b/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)}" diff --git a/test/axon/display_test.exs b/test/axon/display_test.exs new file mode 100644 index 00000000..9dbc63e5 --- /dev/null +++ b/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