Traces execution of the given Axon model with the given inputs, rendering the execution flow as a mermaid flowchart.

You must include kino as a dependency in your project to make use of this function.

Options

• :direction - defines the direction of the graph visual. The value can either be :top_down or :left_right. Defaults to :top_down.

Examples

Given an Axon model:

model = Axon.input("input") |> Axon.dense(32)

You can define input templates for each input:

input = Nx.template({1, 16}, :f32)

And then display the execution flow of the model:

Axon.Display.as_graph(model, input, direction: :top_down)

Traces execution of the given Axon model with the given inputs, rendering the execution flow as a table.

You must include table_rex as a dependency in your project to make use of this function.

Examples

Given an Axon model:

model = Axon.input("input") |> Axon.dense(32)

You can define input templates for each input:

input = Nx.template({1, 16}, :f32)

And then display the execution flow of the model:

Axon.Display.as_table(model, input)