Display.as_table/2 example failing

[bryannaegele]

I can't tell what the exact error is but Display.as_table/2 fails with any model. Narrowing down to the example in the docs generates it.

# Untitled notebook

```elixir
Mix.install([
  {:axon, "~> 0.7"},
  {:nx, "~> 0.9"},
  {:table_rex, "~> 3.1.1"}
])
```

## Section

```elixir
model = Axon.input("input") |> Axon.dense(32)
input = Nx.template({1, 16}, :f32)
Axon.Display.as_table(model, input)
```

** (FunctionClauseError) no function clause matching in Axon.Display.do_axon_to_rows/6    
    
    The following arguments were given to Axon.Display.do_axon_to_rows/6:
    
        # 1
        %Axon.Node{
          id: 5,
          name: #Function<194.34029859/2 in Axon.name/2>,
          mode: :both,
          parent: [4],
          parameters: [
            %Axon.Parameter{
              name: "kernel",
              template: #Function<485.34029859/1 in Axon.shape_fun/2>,
              shape: nil,
              initializer: #Function<3.47872556/3 in Axon.Initializers.glorot_uniform/1>,
              children: nil,
              type: {:f, 32},
              frozen: false,
              kind: :parameter
            },
            %Axon.Parameter{
              name: "bias",
              template: #Function<485.34029859/1 in Axon.shape_fun/2>,
              shape: nil,
              initializer: #Function<23.47872556/2 in Axon.Initializers.zeros/0>,
              children: nil,
              type: {:f, 32},
              frozen: false,
              kind: :parameter
            }
          ],
          args: [:layer, :parameter, :parameter],
          op: :dense,
          policy: #Axon.MixedPrecision.Policy<>,
          hooks: [],
          opts: [],
          global_options: [],
          op_name: :dense,
          meta: %{units: 32, use_bias: true},
          stacktrace: [
            {Axon, :layer, 3, [file: ~c"lib/axon.ex", line: 346]},
            {Axon, :dense, 3, [file: ~c"lib/axon.ex", line: 877]},
            {:elixir, :eval_external_handler, 3, [file: ~c"src/elixir.erl", line: 386]},
            {:erl_eval, :do_apply, 7, [file: ~c"erl_eval.erl", line: 904]},
            {:erl_eval, :expr, 6, [file: ~c"erl_eval.erl", line: 648]},
            {:erl_eval, :exprs, 6, [file: ~c"erl_eval.erl", line: 271]}
          ]
        }
    
        # 2
        %{
          4 => %Axon.Node{
            id: 4,
            name: #Function<195.34029859/2 in Axon.name/2>,
            mode: :both,
            parent: [],
            parameters: [],
            args: [],
            op: :input,
            policy: #Axon.MixedPrecision.Policy<>,
            hooks: [],
            opts: [shape: nil, optional: false],
            global_options: [],
            op_name: :input,
            meta: nil,
            stacktrace: [
              {Axon, :layer, 3, [file: ~c"lib/axon.ex", line: 346]},
              {:elixir, :eval_external_handler, 3, [file: ~c"src/elixir.erl", line: 386]},
              {:erl_eval, :do_apply, 7, [file: ~c"erl_eval.erl", line: 904]},
              {:erl_eval, :expr_list, 7, [file: ~c"erl_eval.erl", line: 1192]},
              {:erl_eval, :expr, 6, [file: ~c"erl_eval.erl", line: 610]},
              {:erl_eval, :expr, 6, [file: ~c"erl_eval.erl", line: 648]}
            ]
          },
          5 => %Axon.Node{
            id: 5,
            name: #Function<194.34029859/2 in Axon.name/2>,
            mode: :both,
            parent: [4],
            parameters: [
              %Axon.Parameter{
                name: "kernel",
                template: #Function<485.34029859/1 in Axon.shape_fun/2>,
                shape: nil,
                initializer: #Function<3.47872556/3 in Axon.Initializers.glorot_uniform/1>,
                children: nil,
                type: {:f, 32},
                frozen: false,
                kind: :parameter
              },
              %Axon.Parameter{
                name: "bias",
                template: #Function<485.34029859/1 in Axon.shape_fun/2>,
                shape: nil,
                initializer: #Function<23.47872556/2 in Axon.Initializers.zeros/0>,
                children: nil,
                type: {:f, 32},
                frozen: false,
                kind: :parameter
              }
            ],
            args: [:layer, :parameter, :parameter],
            op: :dense,
            policy: #Axon.MixedPrecision.Policy<>,
            hooks: [],
            opts: [],
            global_options: [],
            op_name: :dense,
            meta: %{units: 32, use_bias: true},
            stacktrace: [
              {Axon, :layer, 3, [file: ~c"lib/axon.ex", line: 346]},
              {Axon, :dense, 3, [file: ~c"lib/axon.ex", line: 877]},
              {:elixir, :eval_external_handler, 3, [file: ~c"src/elixir.erl", line: 386]},
              {:erl_eval, :do_apply, 7, [file: ~c"erl_eval.erl", line: 904]},
              {:erl_eval, :expr, 6, [file: ~c"erl_eval.erl", line: 648]},
              {:erl_eval, :exprs, 6, [file: ~c"erl_eval.erl", line: 271]}
            ]
          }
        }
    
        # 3
        #Nx.Tensor<
          f32[1][16]
          Nx.TemplateBackend
        >
    
        # 4
        %{}
    
        # 5
        %{}
    
        # 6
        %{num_params: 0, total_param_byte_size: 0}
    
    Attempted function clauses (showing 2 out of 2):
    
        defp do_axon_to_rows(%Axon.Node{id: id, op: structure, op_name: :container, parent: [parents], name: name_fn}, nodes, templates, cache, op_counts, model_info)
        defp do_axon_to_rows(%Axon.Node{
      id: id,
      parent: parents,
      parameters: params,
      name: name_fn,
      opts: opts,
      policy: %{params: {_, bitsize}},
      op_name: op_name
    }, nodes, templates, cache, op_counts, model_info)
    
    (axon 0.7.0) lib/axon/display.ex:93: Axon.Display.do_axon_to_rows/6
    (axon 0.7.0) lib/axon/display.ex:85: Axon.Display.axon_to_rows/6
    (axon 0.7.0) lib/axon/display.ex:39: Axon.Display.as_table/2
    #cell:6xdslgdkoqxorqxd:3: (file)

Pre-0.7 works as expected.

Mix.install([
  {:axon, "~> 0.6.1"},
  {:nx, "~> 0.7"},
  {:exla, "~> 0.7"},
  {:table_rex, "~> 3.1.1"}
])

[heywhy]

The issue was already fixed by a8d3149, the team needs to make a new release.
While waiting for the team you can do {:axon, github: "elixir-nx/axon"} in your mixfile.