Axon

https://hexdocs.pm/axon
Admin
Axon is a deep learning library for Elixir, providing a functional API for building, training, and...
Tokens:199,838
Snippets:1,275
Trust Score:9.7
Update:9 months ago
Show doc for...
Tokens:
Raw
### Install Axon and Dependencies

Source: https://hexdocs.pm/axon/horses_or_humans

Installs the necessary Axon, Nx, EXLA, STB_Image, Req, and Kino libraries for the project. Sets the default backend for Nx to EXLA.

```elixir
Mix.install([
  {:axon, "~> 0.6.0"},
  {:nx, "~> 0.6.0"},
  {:exla, "~> 0.6.0"},
  {:stb_image, "~> 0.6.0"},
  {:req, "~> 0.4.5"},
  {:kino, "~> 0.11.0"}
])

Nx.global_default_backend(EXLA.Backend)
Nx.Defn.global_default_options(compiler: EXLA)
```

--------------------------------

### Install Axon Dependency for Elixir

Source: https://hexdocs.pm/axon/instrumenting_loops_with_metrics

This snippet shows how to install the Axon library using Mix.install, a function for managing dependencies in Elixir. Ensure you have Elixir installed to use this command.

```elixir
Mix.install([
  {:axon, ">= 0.5.0"}
])
```

--------------------------------

### PyTorch Optimizers Examples

Source: https://hexdocs.pm/axon/axon_pytorch

Provides examples of initializing common optimizers in PyTorch, such as Stochastic Gradient Descent (SGD) and Adam, for model parameter updates.

```Python
# Example:
import torch
import torch.optim as optim

optimizer = optim.SGD(model.parameters(), lr=0.01)
optimizer = optim.Adam(model.parameters(), lr=0.001)
```

--------------------------------

### Install Axon Dependencies

Source: https://hexdocs.pm/axon/accelerating_axon

Installs necessary dependencies for Axon, including Axon itself, EXLA, TorchX, Benchee, and Kino. Kino is overridden to ensure a specific version is used.

```elixir
Mix.install([
  {:axon, ">= 0.5.0"},
  {:exla, ">= 0.5.0"},
  {:torchx, ">= 0.5.0"},
  {:benchee, "~> 1.1"},
  {:kino, ">= 0.9.0", override: true}
])
```

--------------------------------

### Axon Variance Scaling Initializer Example

Source: https://hexdocs.pm/axon/Axon.Initializers

Demonstrates initializing parameters with variance scaling using Axon.Initializers.variance_scaling. Examples include default settings and custom modes and distributions for various tensor shapes and types.

```Elixir
iex> init_fn = Axon.Initializers.variance_scaling()
iex> t = init_fn.({2, 2}, {:f, 32}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:f, 32}
```

```Elixir
iex> init_fn = Axon.Initializers.variance_scaling(mode: :fan_out, distribution: :truncated_normal)
iex> t = init_fn.({2, 2}, {:bf, 16}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:bf, 16}
```

```Elixir
iex> init_fn = Axon.Initializers.variance_scaling(mode: :fan_out, distribution: :normal)
iex> t = init_fn.({64, 3, 32, 32}, {:f, 32}, Nx.Random.key(1))
iex> Nx.shape(t)
{64, 3, 32, 32}
iex> Nx.type(t)
{:f, 32}
```

--------------------------------

### Install Axon and related dependencies

Source: https://hexdocs.pm/axon/mnist

Installs the necessary Axon, Nx, Exla, and Req libraries using Mix.install. These are essential for building and running neural networks with Axon.

```elixir
Mix.install([
  {:axon, "~> 0.3.0"},
  {:nx, "~> 0.4.0", override: true},
  {:exla, "~> 0.4.0"},
  {:req, "~> 0.3.1"}
])
```

--------------------------------

### Axon Initializers: Glorot Normal Initialization Examples

Source: https://hexdocs.pm/axon/Axon.Initializers

Provides examples of using the `glorot_normal` initializer, also known as Xavier normal. It demonstrates initializing parameters with the Glorot normal distribution and how to specify a scale for the output distribution.

```elixir
iex> init_fn = Axon.Initializers.glorot_normal()
iex> t = init_fn.({2, 2}, {:f, 32}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:f, 32}
```

```elixir
iex> init_fn = Axon.Initializers.glorot_normal(scale: 1.0e-3)
iex> t = init_fn.({2, 2}, {:bf, 16}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:bf, 16}
```

--------------------------------

### Install Axon and Kino Libraries

Source: https://hexdocs.pm/axon/complex_models

This snippet shows how to install the necessary Axon and Kino libraries using Mix.install. These are required for building and displaying Axon models.

```elixir
Mix.install([
  {:axon, "~> 0.5.0"},
  {:kino, "~> 0.9.0"}
])
```

--------------------------------

### AxonInterp: Inference Example

Source: https://hexdocs.pm/axon/_interp/AxonInterp

An example demonstrating the workflow of setting input tensor, invoking inference, and retrieving the output tensor using AxonInterp functions.

```elixir
output_bin0 = session()
  |> AxonInterp.set_input_tensor(0, input_bin0)
  |> AxonInterp.invoke()
  |> AxonInterp.get_output_tensor(0)
```

--------------------------------

### Axon Initializers: Glorot Uniform Initialization Examples

Source: https://hexdocs.pm/axon/Axon.Initializers

Shows how to use the `glorot_uniform` initializer, also known as Xavier uniform. Examples include default usage and specifying a custom scale for the output distribution, demonstrating different tensor shapes and types.

```elixir
iex> init_fn = Axon.Initializers.glorot_uniform()
iex> t = init_fn.({2, 2}, {:f, 32}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:f, 32}
```

```elixir
iex> init_fn = Axon.Initializers.glorot_uniform(scale: 1.0e-3)
iex> t = init_fn.({2, 2}, {:bf, 16}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:bf, 16}
```

--------------------------------

### Axon Initializers: Full Value Initialization Example

Source: https://hexdocs.pm/axon/Axon.Initializers

Illustrates the use of the `full` initializer to create a tensor where all elements are set to a specified value. The example shows how to create an initialization function and then apply it to a given shape and type.

```elixir
iex> init_fn = Axon.Initializers.full(1.00)
iex> out = init_fn.({2, 2}, {:f, 32})
iex> out
#Nx.Tensor<
  f32[2][2]
  [
    [1.0, 1.0],
    [1.0, 1.0]
  ]
>
```

--------------------------------

### Install Axon and Kino Dependencies - Elixir

Source: https://hexdocs.pm/axon/custom_layers

This snippet demonstrates how to install the Axon and Kino libraries using Mix.install, which is common in Elixir projects for managing dependencies. Ensure you have Elixir and Mix installed to run this command.

```elixir
Mix.install([
  {:axon, ">= 0.5.0"},
  {:kino, ">= 0.9.0"}
])
```

--------------------------------

### Install Axon and Nx Libraries

Source: https://hexdocs.pm/axon/custom_models_loss_optimizers

This snippet shows how to install the Axon and Nx Elixir libraries using Mix.install. It specifies the GitHub repositories and versions for each dependency. This is typically used at the beginning of a project to set up the development environment.

```elixir
Mix.install([
  {:axon, github: "elixir-nx/axon"},
  {:nx, "~> 0.3.0", github: "elixir-nx/nx", sparse: "nx", override: true}
])
```

--------------------------------

### Install Axon and Dependencies

Source: https://hexdocs.pm/axon/metric-learning

Installs necessary Elixir dependencies for Axon, Scidata, EXLA, STBImage, and Kino Vega Lite. Sets EXLA as the global default backend for Nx operations.

```elixir
Mix.install([
  {:axon, "~> 0.7"},
  {:scidata, "~> 0.1.9"},
  {:exla, "~> 0.9"},
  {:stb_image, "~> 0.5.2"},
  {:kino_vega_lite, "~>0.1.13"}
])

Nx.global_default_backend(EXLA.Backend)
Nx.Defn.global_default_options(compiler: EXLA)
```

--------------------------------

### Install Axon Library using Mix

Source: https://hexdocs.pm/axon/using_loop_event_handlers

This snippet shows how to install the Axon library using Mix. It requires Elixir version 0.5.0 or higher. The installation is a prerequisite for using Axon in your project.

```elixir
Mix.install([
  {:axon, ">= 0.5.0"}
])copy
```

--------------------------------

### Supervised Learning Batch Step Example

Source: https://hexdocs.pm/axon/Axon.Loop

Illustrates a `batch_step` implementation for supervised learning, showing how to handle inputs and targets, compute gradients, and update model parameters.

```elixir
defn batch_step({inputs, targets}, state) do
  %{parameters: params, optimizer_state: optim_state} = state

  gradients = grad(params, objective_fn.(&1, inputs, targets))
  {updates, new_optim_state} = optimizer.(optim_state, params, gradients)

  new_params = apply_updates(params, updates)

  %{parameters: new_params, optimizer_state: optim_state}
end
```

--------------------------------

### Install Axon and Dependencies for Elixir

Source: https://hexdocs.pm/axon/xor

Installs necessary Elixir libraries for neural network operations, including Axon, Nx, EXLA, and Kino Vega Lite. These are essential for building and training neural networks.

```elixir
Mix.install([
  {:axon, "~> 0.3.0"},
  {:nx, "~> 0.4.0", override: true},
  {:exla, "~> 0.4.0"},
  {:kino_vega_lite, "~> 0.1.6"}
])

Nx.Defn.default_options(compiler: EXLA)

alias VegaLite, as: Vlcopy
```

--------------------------------

### Install Axon and Dependencies

Source: https://hexdocs.pm/axon/credit_card_fraud

Installs necessary Elixir packages for machine learning, including Axon, NX, EXLA, Explorer, and Kino. It also configures the default backend for NX to EXLA.Backend.

```elixir
Mix.install([
  {:axon, "~> 0.3.0"},
  {:nx, "~> 0.4.0", override: true},
  {:exla, "~> 0.4.0"},
  {:explorer, "~> 0.3.1"},
  {:kino, "~> 0.7.0"}
])

Nx.Defn.default_options(compiler: EXLA)
Nx.global_default_backend(EXLA.Backend)
```

--------------------------------

### Install Axon Dependencies

Source: https://hexdocs.pm/axon/lstm_generation

Installs necessary Elixir dependencies for the Axon project, including Axon, Nx, Exla, and Req. It also configures Nx to use EXLA as the default backend and compiler.

```elixir
Mix.install([
  {:axon, "~> 0.3.0"},
  {:nx, "~> 0.4.0", override: true},
  {:exla, "~> 0.4.0"},
  {:req, "~> 0.3.1"}
])

Nx.Defn.default_options(compiler: EXLA)
Nx.global_default_backend(EXLA.Backend)
```

--------------------------------

### Install Axon and Dependencies

Source: https://hexdocs.pm/axon/onnx_to_axon

Installs necessary Elixir dependencies for working with Axon, EXLA, Axon_ONNX, and image/network libraries. It also shows how to configure the XLA target for GPU acceleration if available, by setting the `XLA_TARGET` environment variable.

```elixir
Mix.install(
  [
    {:axon, "~> 0.5.0"},
    {:exla, "~> 0.5.0"},
    {:axon_onnx, "~> 0.4.0"},
    {:stb_image, "~> 0.6.0"},
    {:kino, "~> 0.9.0"},
    {:req, "~> 0.3.8"}
  ]
  # for Nvidia GPU change to "cuda111" for CUDA 11.1+ or "cuda118" for CUDA 11.8
  # CUDA 12.x not supported by XLA
  # or you can put this value in ENV variables in Livebook settings
  # XLA_TARGET=cuda111
  # system_env: %{"XLA_TARGET" => xla_target}
)

```

--------------------------------

### Install Dependencies for Axon Project

Source: https://hexdocs.pm/axon/fashionmnist_vae

Installs necessary Elixir dependencies including Exla, Nx, Axon, Req, Kino, Scidata, StbImage, Kino_VegaLite, VegaLite, and TableRex. It also sets EXLA as the global default backend for Nx operations.

```elixir
Mix.install([
  {:exla, "~> 0.4.0"},
  {:nx, "~> 0.4.0", override: true},
  {:axon, "~> 0.3.0"},
  {:req, "~> 0.3.1"},
  {:kino, "~> 0.7.0"},
  {:scidata, "~> 0.1.9"},
  {:stb_image, "~> 0.5.2"},
  {:kino_vega_lite, "~> 0.1.6"},
  {:vega_lite, "~> 0.1.6"},
  {:table_rex, "~> 3.1.1"}
])

alias VegaLite, as: Vl

# This speeds up all our `Nx` operations without having to use `defn`
Nx.global_default_backend(EXLA.Backend)
```

--------------------------------

### Install Axon and Dependencies

Source: https://hexdocs.pm/axon/fashionmnist_autoencoder

Installs the necessary Axon, Nx, EXLA, and Scidata libraries for building and training the autoencoder. Ensures compatibility by overriding Nx version and setting EXLA as the default compiler.

```elixir
Mix.install([
  {:axon, "~> 0.3.0"},
  {:nx, "~> 0.4.0", override: true},
  {:exla, "~> 0.4.0"},
  {:scidata, "~> 0.1.9"}
])

Nx.Defn.default_options(compiler: EXLA)
```

--------------------------------

### Sigmoid Activation Example (Elixir)

Source: https://hexdocs.pm/axon/Axon.Activations

Provides a conceptual overview of the Sigmoid activation function in Axon.Activations. Although no direct code example is shown, it explains the mathematical formula and notes its implementation detail regarding cached logits for numerical stability in cross-entropy calculations.

```elixir
# Sigmoid activation.
f(xi)=11+e−xi f(x_i) = \frac{1}{1 + e^{-x_i}} f(xi)=1+e−xi1
**Implementation Note: Sigmoid logits are cached as metadata in the expression and can be used in calculations later on. For example, they are used in cross-entropy calculations for better stability.**
```

--------------------------------

### ELU Activation Example

Source: https://hexdocs.pm/axon/Axon.Activations

Demonstrates the usage of the Exponential linear unit (ELU) activation function, highlighting its equivalence to CELU with alpha=1. Examples cover different input tensors and data types.

```elixir
iex> Axon.Activations.elu(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0]))
#Nx.Tensor<
  f32[7]
  [-0.9502129554748535, -0.8646647334098816, -0.6321205496788025, 0.0, 1.0, 2.0, 3.0]
>

iex> Axon.Activations.elu(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}))
#Nx.Tensor<
  bf16[2][3]
  [
    [-0.62890625, -0.86328125, -0.94921875],
    [1.0, 2.0, 3.0]
  ]
>copy
```

--------------------------------

### Mish Activation Example (Elixir)

Source: https://hexdocs.pm/axon/Axon.Activations

Demonstrates the Mish activation function within Axon.Activations, using Nx Tensors. The examples cover a 1D tensor and a 2D tensor, showing the output for f32 and bf16 types respectively.

```elixir
iex> Axon.Activations.mish(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], type: {:f, 32}, names: [:data]))
#Nx.Tensor<
  f32[data: 7]
  [-0.14564745128154755, -0.2525014877319336, -0.30340147018432617, 0.0, 0.8650984168052673, 1.9439589977264404, 2.98653507232666]
>

iex> Axon.Activations.mish(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 3]
  [
    [-0.30078125, -0.25, -0.1435546875],
    [0.86328125, 1.9375, 2.96875]
  ]
>
```

--------------------------------

### SELU Activation Example (Elixir)

Source: https://hexdocs.pm/axon/Axon.Activations

Illustrates the Scaled Exponential Linear Unit (SELU) activation function from Axon.Activations using Nx Tensors. The examples show its application on a 1D tensor and a 2D tensor with bf16 type, demonstrating its self-normalizing properties.

```elixir
iex> Axon.Activations.selu(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], names: [:data]))
#Nx.Tensor<
  f32[data: 7]
  [-1.670568823814392, -1.5201665163040161, -1.1113307476043701, 0.0, 1.0507010221481323, 2.1014020442962646, 3.1521029472351074]
>

iex> Axon.Activations.selu(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 3]
  [
    [-1.09375, -1.5, -1.65625],
    [1.046875, 2.09375, 3.140625]
  ]
>
```

--------------------------------

### CELU Activation Example

Source: https://hexdocs.pm/axon/Axon.Activations

Provides examples of using the Continuously-differentiable exponential linear unit (CELU) activation function with different input tensors and data types. It also shows an error case when the alpha parameter is zero.

```elixir
iex> Axon.Activations.celu(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0]))
#Nx.Tensor<
  f32[7]
  [-0.9502129554748535, -0.8646647334098816, -0.6321205496788025, 0.0, 1.0, 2.0, 3.0]
>

iex> Axon.Activations.celu(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}))
#Nx.Tensor<
  bf16[2][3]
  [
    [-0.62890625, -0.86328125, -0.94921875],
    [1.0, 2.0, 3.0]
  ]
>copy

iex> Axon.Activations.celu(Nx.tensor([0.0, 1.0, 2.0], type: {:f, 32}), alpha: 0.0)
** (ArgumentError) :alpha must be non-zero in CELU activationcopy
```

--------------------------------

### One-Hot Encoding Example (Nx/Elixir)

Source: https://hexdocs.pm/axon/lstm_generation

Demonstrates how to perform one-hot encoding using Nx in Elixir. This is crucial for preparing categorical target data for neural networks, allowing the model to output probabilities for each category. The example shows creating a tensor and comparing it with an identity matrix to achieve the one-hot representation.

```Elixir
Nx.tensor([
  [0],
  [1],
  [2]
])
|> Nx.equal(Nx.iota({1, 3}))
```

--------------------------------

### Start and Inspect YOLOv4 Model

Source: https://hexdocs.pm/axon/_interp/yolov4

These snippets demonstrate how to start the YOLOv4 inference process and retrieve information about the loaded model. `YOLOv4.start_link([])` initiates the model's background process, while `AxonInterp.info(YOLOv4)` displays the model's properties, including input/output details and configuration.

```elixir
# AxonInterp.stop(Resnet18)
YOLOv4.start_link([])
```

```elixir
AxonInterp.info(YOLOv4)
```

--------------------------------

### Log-Sum-Exp Activation Example (Elixir)

Source: https://hexdocs.pm/axon/Axon.Activations

Illustrates the log_sumexp activation function from Axon.Activations using Nx Tensors. The examples show its application on a 1D tensor and a 2D tensor with bf16 type, producing a single output value per batch.

```elixir
iex> Axon.Activations.log_sumexp(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], names: [:data]))
#Nx.Tensor<
  f32[data: 1]
  [3.4577627182006836]
>

iex> Axon.Activations.log_sumexp(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 1]
  [
    [-0.59375],
    [3.390625]
  ]
>
```

--------------------------------

### Run Axon Training with Custom Metrics

Source: https://hexdocs.pm/axon/writing_custom_metrics

This Elixir snippet shows how to set up dummy training data using `Stream.repeatedly` and then run the Axon training loop with the defined model and custom metrics. The `Axon.Loop.run/3` function starts the training process for a specified number of iterations. The output shows an example of how the custom metric ('my weird metric') is reported alongside the loss.

```elixir
train_data = 
  Stream.repeatedly(fn ->
    {xs, _next_key} =
      :random.uniform(9999)
      |> Nx.Random.key()
      |> Nx.Random.normal(shape: {8, 1})

    ys = Nx.sin(xs)
    {xs, ys}
  end)

Axon.Loop.run(loop, train_data, %{}, iterations: 1000)
```

--------------------------------

### Axon Uniform Initializer Example

Source: https://hexdocs.pm/axon/Axon.Initializers

Illustrates initializing parameters with a random uniform distribution using Axon.Initializers.uniform. It covers default usage and custom scale with different data types.

```Elixir
iex> init_fn = Axon.Initializers.uniform()
iex> t = init_fn.({2, 2}, {:f, 32}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:f, 32}
```

```Elixir
iex> init_fn = Axon.Initializers.uniform(scale: 1.0e-3)
iex> t = init_fn.({2, 2}, {:bf, 16}, Nx.Random.key(1))
iex> Nx.shape(t)
{2, 2}
iex> Nx.type(t)
{:bf, 16}
```

--------------------------------

### Axon Zeros Initializer Example

Source: https://hexdocs.pm/axon/Axon.Initializers

Shows how to initialize parameters to zero using Axon.Initializers.zeros. It demonstrates creating a tensor with a specified shape and data type filled with zeros.

```Elixir
iex> init_fn = Axon.Initializers.zeros()
iex> out = init_fn.({2, 2}, {:f, 32})
iex> out
#Nx.Tensor<
  f32[2][2]
  [
    [0.0, 0.0],
    [0.0, 0.0]
  ]
>
```

--------------------------------

### Output Transform Function Example

Source: https://hexdocs.pm/axon/Axon.Loop

Provides an example of an `output_transform` function used with `Axon.Loop.trainer/4` to extract specific values, such as the trained model state, from the final loop state.

```elixir
output_transform = fn state ->
  state.step_state[:model_state]
end
```

--------------------------------

### Axon Orthogonal Initializer Example

Source: https://hexdocs.pm/axon/Axon.Initializers

Demonstrates initializing parameters with a random orthogonal distribution using Axon.Initializers.orthogonal. It shows how to specify tensor shape, data type, and random key.

```Elixir
iex> init_fn = Axon.Initializers.orthogonal()
iex> t = init_fn.({3, 3}, {:f, 32}, Nx.Random.key(1))
iex> Nx.type(t)
{:f, 32}
iex> Nx.shape(t)
{3, 3}
```

```Elixir
iex> init_fn = Axon.Initializers.orthogonal()
iex> t = init_fn.({1, 2, 3, 4}, {:f, 64}, Nx.Random.key(1))
iex> Nx.type(t)
{:f, 64}
iex> Nx.shape(t)
{1, 2, 3, 4}
```

--------------------------------

### Run AxonInterp YOLOv4 Demo

Source: https://hexdocs.pm/axon/_interp/readme

Instructions for running the Object Detection: YOLOv4 demo provided with AxonInterp. This involves downloading a pre-trained model, navigating to the demo directory, fetching dependencies, and executing the demo script.

```shell
cd demo
mix deps.get
mix run -e "DemoCandy.run"
```

--------------------------------

### Start and Inspect Resnet18 Module

Source: https://hexdocs.pm/axon/_interp/resnet18

Provides Elixir code snippets to start the Resnet18 inference module linked as a process and to display its properties using AxonInterp.info.

```elixir
# AxonInterp.stop(Resnet18)
Resnet18.start_link([])
```

```elixir
AxonInterp.info(Resnet18)
```

--------------------------------

### Install Dependencies for AxonInterp

Source: https://hexdocs.pm/axon/_interp/resnet18

Installs necessary Elixir dependencies for AxonInterp, CImg, Nx, and Kino. It also configures Nx to use EXLA for default definition options.

```elixir
Mix.install(
  [
    {:axon_interp, "~> 0.1.0"},
    {:cimg, "~> 0.1.14"},
    {:nx, "~> 0.4.0"},
    {:kino, "~> 0.7.0"}
  ],
  config: [
    nx: [default_defn_options: [compiler: EXLA]]
  ]
  # system_env: [{"NNINTERP", "Axon"}]
)
```

--------------------------------

### Start Nx.Serving in Application Supervision Tree

Source: https://hexdocs.pm/axon/Axon

Demonstrates how to include an Nx.Serving component in an Elixir application's supervision tree. This ensures the serving process is managed and restarted as needed. It specifies the serving configuration, a name for the serving process, and a batch timeout.

```elixir
children = [
  ...,
  {Nx.Serving, serving: build_serving(), name: MyApp.Serving, batch_timeout: 100}
]
```

--------------------------------

### Max Pooling Example - Elixir

Source: https://hexdocs.pm/axon/Axon.Layers

Provides an example of using Axon.Layers.max_pool for standard max pooling operations. Max pooling is commonly used for downsampling feature maps after convolutional layers.

```elixir
# View Source
# max_pool(input, opts \ [])
# __ View Source
# Functional implementation of a general dimensional max pooling layer.
# Pooling is applied to the spatial dimension of the input tensor. Max pooling returns the maximum element in each valid window of the input tensor. It is often used after convolutional layers to downsample the input even further.

```

--------------------------------

### Build and Run Axon Model with Default Backend

Source: https://hexdocs.pm/axon/accelerating_axon

Builds an Axon model and executes it using the currently configured default Nx backend. This demonstrates a typical workflow after setting the backend.

```elixir
{inputs, _next_key} =
  Nx.Random.key(9999)
  |> Nx.Random.uniform(shape: {2, 128})

{init_fn, predict_fn} = Axon.build(model)
params = init_fn.(inputs, %{})
predict_fn.(params, inputs)
```

--------------------------------

### Initialization Function Example

Source: https://hexdocs.pm/axon/Axon.Loop

Demonstrates the `init` function, an arity-0 function responsible for providing the initial step state for the Axon loop, often including initial model parameters.

```elixir
init = fn ->
  %{params: Axon.init(model)}
end
```

--------------------------------

### Define and Train an Axon Neural Network

Source: https://hexdocs.pm/axon/Axon

Provides a complete example of defining a neural network using Axon layers, from input to output. It then demonstrates how to set up a training loop using Axon.Loop.trainer and Axon.Loop.run for model optimization.

```elixir
model =
  Axon.input("input_0", shape: {nil, 784})
  |> Axon.dense(128, activation: :relu)
  |> Axon.layer_norm()
  |> Axon.dropout()
  |> Axon.dense(10, activation: :softmax)

IO.inspect model

model_state =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(learning_rate: 0.005))
  |> Axon.Loop.run(train_data, epochs: 10, compiler: EXLA)
```

--------------------------------

### Axon linear Activation Function Example

Source: https://hexdocs.pm/axon/Axon.Activations

Illustrates the basic linear activation function (identity) using Axon.Activations.linear. Examples are provided for f32 and bf16 tensor types, showing that the output is identical to the input.

```iex
iex> Axon.Activations.linear(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], names: [:data]))
#Nx.Tensor<
  f32[data: 7]
  [-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0]
>
```

```iex
iex> Axon.Activations.linear(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 3]
  [
    [-1.0, -2.0, -3.0],
    [1.0, 2.0, 3.0]
  ]
>
```

--------------------------------

### Build and Execute Axon Model

Source: https://hexdocs.pm/axon/Axon

Explains the process of building an Axon model using Axon.build/2, which returns initialization and prediction functions. These functions are then used to initialize model parameters and make predictions with given inputs.

```elixir
{init_fn, predict_fn} = Axon.build(model)

params = init_fn.(Nx.template({1, 1}, {:f, 32}), %{})
predict_fn.(params, inputs)
```

--------------------------------

### Axon leaky_relu Activation Function Example

Source: https://hexdocs.pm/axon/Axon.Activations

Demonstrates the leaky rectified linear unit activation function (leaky_relu) using Axon.Activations.leaky_relu. Examples show its application with a custom alpha value on f32 tensors.

```iex
iex> Axon.Activations.leaky_relu(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], names: [:data]), alpha: 0.5)
#Nx.Tensor<
  f32[data: 7]
  [-1.5, -1.0, -0.5, 0.0, 1.0, 2.0, 3.0]
>
```

```iex
iex> Axon.Activations.leaky_relu(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], names: [:batch, :data]), alpha: 0.5)
#Nx.Tensor<
  f32[batch: 2][data: 3]
  [
    [-0.5, -1.0, -1.5],
    [1.0, 2.0, 3.0]
  ]
>
```

--------------------------------

### Axon Training Loop with Visualization and Early Stopping

Source: https://hexdocs.pm/axon/fashionmnist_vae

Sets up and runs an Axon training loop. It includes early stopping, iteration completion handling to render example predictions, validation, loss plotting, and uses EXLA compiler for optimization. Assumes `model`, `test_data`, `train_data`, and `Data` module are defined.

```elixir
combined_input_output = fn params, image_index ->
  test_image = test_images[[images: image_index]]
  reconstructed_image = Axon.predict(model, params, test_image) |> Nx.squeeze(axes: [0])
  Nx.concatenate([test_image, reconstructed_image], axis: :width)
end

frame = Kino.Frame.new() |> Kino.render()

render_example_handler = fn state ->
  # state.step_state[:model_state] contains the model params when this event is fired
  params = state.step_state[:model_state]
  image_index = Enum.random(0..(Nx.axis_size(test_images, :images) - 1))
  image = combined_input_output.(params, image_index) |> Data.image_to_kino(200, 400)
  Kino.Frame.render(frame, image)
  Kino.Frame.append(frame, "Epoch: #{state.epoch}, Iteration: #{state.iteration}")
  {:continue, state}
end

params =
  model
  |> Axon.Loop.trainer(:mean_squared_error, Polaris.Optimizers.adamw(learning_rate: 0.001))
  |> KinoAxon.kino_early_stop()
  |> Axon.Loop.handle(:iteration_completed, render_example_handler, every: 450)
  |> Axon.Loop.validate(model, test_data)
  |> KinoAxon.plot_losses()
  |> Axon.Loop.run(train_data, %{}, epochs: 40, compiler: EXLA)

:ok
```

--------------------------------

### Axon hard_sigmoid Activation Function Example

Source: https://hexdocs.pm/axon/Axon.Activations

Shows the application of the hard sigmoid activation function via Axon.Activations.hard_sigmoid. Examples are provided for f32 and bf16 tensor types, illustrating the bounded output between 0 and 1.

```iex
iex> Axon.Activations.hard_sigmoid(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], names: [:data]))
#Nx.Tensor<
  f32[data: 7]
  [0.0, 0.0, 0.0, 0.20000000298023224, 0.4000000059604645, 0.6000000238418579, 0.800000011920929]
>
```

```iex
iex> Axon.Activations.hard_sigmoid(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 3]
  [
    [0.0, 0.0, 0.0],
    [0.3984375, 0.59765625, 0.796875]
  ]
>
```

--------------------------------

### Build and Initialize Axon Model Parameters

Source: https://hexdocs.pm/axon/sequential_models

This code shows how to build an Axon model into executable functions and initialize its trainable parameters. `Axon.build` compiles the model into `init_fn` and `predict_fn`. The `init_fn` is then used with an input template to create an initial parameter map, which contains the weights and biases for each layer.

```Elixir
{init_fn, predict_fn} = Axon.build(model)
params = init_fn.(template, %{})
```

--------------------------------

### Benchmark Elixir vs. EXLA JIT Compiled Init Functions

Source: https://hexdocs.pm/axon/accelerating_axon

This benchmark compares the performance of a standard Elixir initialization function against an EXLA JIT-compiled version. It uses the Benchee library to measure execution time and memory usage. Ensure `init_fn` and `inputs` are defined beforehand.

```elixir
Benchee.run(
  %{
    "elixir init" => fn -> init_fn.(inputs, %{}) end,
    "exla init" => fn -> exla_init_fn.(inputs, %{}) end
  },
  time: 10,
  memory_time: 5,
  warmup: 2
)
```

--------------------------------

### Build and Run an Axon Model

Source: https://hexdocs.pm/axon/complex_models

Shows how to build an Axon model using Axon.build/2, which returns an initialization function and a prediction function. It then demonstrates how to initialize the model parameters and make a prediction using the predict function with sample input data.

```elixir
{init_fn, predict_fn} = Axon.build(out)

params = init_fn.(template, %{})
predict_fn.(params, Nx.iota({2, 8}, type: :f32))
```

--------------------------------

### Axon hard_tanh Activation Function Example

Source: https://hexdocs.pm/axon/Axon.Activations

Provides examples for the hard hyperbolic tangent activation function (hard_tanh) using Axon.Activations.hard_tanh. The function is applied to f32 and bf16 tensors, demonstrating its clamping behavior between -1 and 1.

```iex
iex> Axon.Activations.hard_tanh(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], names: [:data]))
#Nx.Tensor<
  f32[data: 7]
  [-1.0, -1.0, -1.0, 0.0, 1.0, 1.0, 1.0]
>
```

```iex
iex> Axon.Activations.hard_tanh(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 3]
  [
    [-1.0, -1.0, -1.0],
    [1.0, 1.0, 1.0]
  ]
>
```

--------------------------------

### Axon hard_silu Activation Function Example

Source: https://hexdocs.pm/axon/Axon.Activations

Demonstrates the hard sigmoid weighted linear unit activation function (hard_silu) using Axon.Activations.hard_silu. Examples are shown for f32 and bf16 tensor types, illustrating its piecewise definition.

```iex
iex> Axon.Activations.hard_silu(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], names: [:data]))
#Nx.Tensor<
  f32[data: 7]
  [-0.0, -0.0, -0.0, 0.0, 0.4000000059604645, 1.2000000476837158, 2.4000000953674316]
>
```

```iex
iex> Axon.Activations.hard_silu(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 3]
  [
    [-0.0, -0.0, -0.0],
    [0.3984375, 1.1953125, 2.390625]
  ]
>
```

--------------------------------

### Axon Basic Training Loop Examples

Source: https://hexdocs.pm/axon/axon_pytorch

Demonstrates training loops in Axon. The `Axon.Loop` module offers a high-level approach, while a manual loop structure involving gradient calculation and parameter updates is also shown.

```Elixir
# High-level approach using Axon.Loop:
model_state = Axon.Loop.trainer(model, :categorical_cross_entropy, optimizer)
           |> Axon.Loop.run(train_data, epochs: 10, compiler: EXLA)

# Manual loop structure (simplified):
model = # ...
loss_fn = &Axon.Losses.categorical_cross_entropy/2
optimizer = Polaris.Optimizers.adam(learning_rate: 0.001)
{init_fn, predict_fn} = Axon.build(model, compiler: EXLA, mode: :train)
params = init_fn.(input_template, key)
opt_state = Polaris.Optimizers.init(optimizer, params)

for epoch <- 1..10 do
  Enum.reduce(train_data, {params, opt_state}, fn {inputs, targets}, {params, opt_state} ->
    # Define a gradient function
    {{_preds, loss}, grads} =
      Nx.Defn.value_and_grad(
        inputs,
        fn inputs ->
          preds = predict_fn.(params, inputs)
          loss = loss_fn.(targets, preds)

          {preds, loss}
        end,
        fn {_preds, loss} -> loss end
      )

    {updates, opt_state} = Polaris.Optimizers.update(optimizer, grads, params, opt_state)
    params = Polaris.Updates.apply_updates(params, updates)
    {params, opt_state}
  end)
end
```

--------------------------------

### Axon Initializers: Zeros Initialization Example

Source: https://hexdocs.pm/axon/Axon.Initializers

Demonstrates how to use the `zeros` initializer function from Axon.Initializers to create a tensor filled with zeros. This function returns an initialization function that takes shape and type arguments and returns a tensor.

```elixir
init_fn = Axon.Initializers.zeros()
init_fn.({1, 2}, {:f, 32})
```

--------------------------------

### Build Axon Model Initialization and Prediction Functions

Source: https://hexdocs.pm/axon/your_first_axon_model

This snippet uses Axon.build/2 to transform an Axon model definition into separate initialization and prediction functions. These functions are essential for preparing the model's parameters and for running inference.

```elixir
{init_fn, predict_fn} = Axon.build(input)
```

--------------------------------

### ReLU6 Activation Example (Elixir)

Source: https://hexdocs.pm/axon/Axon.Activations

Presents examples of the ReLU6 activation function from Axon.Activations using Nx Tensors. It shows the function's behavior on a 1D tensor and a 2D tensor with bf16 type, capping the output at 6.

```elixir
iex> Axon.Activations.relu6(Nx.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0]))
#Nx.Tensor<
  f32[7]
  [0.0, 0.0, 0.0, 0.0, 1.0, 2.0, 3.0]
>

iex> Axon.Activations.relu6(Nx.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]], type: {:bf, 16}, names: [:batch, :data]))
#Nx.Tensor<
  bf16[batch: 2][data: 3]
  [
    [0.0, 0.0, 0.0],
    [1.0, 2.0, 3.0]
  ]
>
```

--------------------------------

### Prepare and Visualize Training Data

Source: https://hexdocs.pm/axon/fashionmnist_vae

Prepares training and testing data by batching images and then visualizes the first input and target batches using Kino. This function assumes `train_images`, `test_images`, and `Data.prepare_training_data` are defined elsewhere.

```elixir
batch_size = 128

train_data = Data.prepare_training_data(train_images, 128)
test_data = Data.prepare_training_data(test_images, 128)

{input_batch, target_batch} = Enum.at(train_data, 0)
Kino.render(input_batch[[images: 0]] |> Data.image_to_kino())
Kino.render(target_batch[[images: 0]] |> Data.image_to_kino())
```
Tokens:
Raw