View Source Axon.Display (Axon v0.3.1)

Module for rendering various visual representations of Axon models.

Link to this section Summary

Functions

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

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

Link to this section Functions

Link to this function

as_graph(axon, input_templates, opts \\ [])

View Source

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

Options

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

examples

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)
Link to this function

as_table(axon, input_templates)

View Source

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

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)