Axon: Neural Networks for Elixir, Built on Nx

What Axon Solves and Who Uses It

Axon addresses the lack of a first-class neural network library in the Elixir ecosystem. It sits on top of Nx (Numerical Elixir), a tensor computation library that provides hardware acceleration through pluggable compiler backends. Before Axon, Elixir developers who wanted to train or run neural networks had to call out to Python processes or use raw matrix operations.

The target user is an Elixir developer who wants to build, train, or deploy neural networks without leaving the BEAM runtime. This includes teams building real-time data pipelines in Elixir who want inference as a first-class operation in the same process, and Elixir engineers who want ML without a Python dependency.

Axon is not a general-purpose ML framework for Python developers, and it is not a zero-configuration tool. The README notes that practical workloads require including an Nx compiler such as EXLA, meaning Axon alone is incomplete for production use.

The Three-Layer API Architecture

The README describes three distinct API layers, each independently usable.

The functional API is the lowest level. It implements numerical definitions (`defn`) for activations, initializers, layers, losses, and metrics. Because these are numerical definitions, they can run on any Nx backend and be composed with other numerical definitions. This includes modules like `Axon.Activations`, `Axon.Layers`, `Axon.Losses`, and `Axon.Metrics`.

The model creation API builds on the functional layer. It manages model initialization and application as two separate concerns, which Axon surfaces as two functions from `Axon.build/2`. A model is an Elixir struct; Axon's `Axon.Display` module can render a table summary showing layer names, shapes, and parameter counts. The design choice is that the model struct is only concerned with initialization and prediction, not training.

The training API provides convenience wrappers and event handlers for training loops, using an approach the README describes as inspired by PyTorch Ignite. The general pattern is: define a model, create a loop using a factory like `Axon.Loop.trainer/3`, attach metrics and event handlers, then run the loop on data.

Installing Axon and Running a First Model

Axon requires Elixir and the Mix build tool. Add Axon to your project dependencies in mix.exs:

def deps do
  [
    {:axon, "~> 0.6"}
  ]
end

For any practical deep learning workload, include an Nx compiler as well:

def deps do
  [
    {:axon, "~> 0.6"},
    {:exla, "~> 0.6"},
  ]
end

A basic model pipes input through dense layers using the `|>` operator:

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

{init_fn, predict_fn} = Axon.build(model, compiler: EXLA)
params = init_fn.(Nx.template({1, 784}, :f32), %{})
predict_fn.(params, input)

The `Axon.build/2` call returns an initialisation function and a prediction function. The initialisation function takes an input template and an empty parameter map and returns the initialised parameters. Separating initialisation from prediction is a deliberate design choice noted in the README: the model struct is not responsible for training state.

Training a Model With the Loop API

The training API follows a four-step pattern: define a model, define a loop using a factory method, attach metrics and event handlers, and run the loop on data.

model =
  Axon.input("input", shape: {nil, 784})
  |> Axon.dense(128)
  |> Axon.dense(10, activation: :softmax)

%Axon.Loop.State{step_state: %{model_state: model_state}} =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(0.005))
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.handle(:iteration_completed, &log_metrics/1, every: 50)
  |> Axon.Loop.run(data, %{}, epochs: 10, compiler: EXLA)

Axon uses Polaris for its optimisation API. The README notes an important design point: the optimisation API does not depend on Axon models directly. You can use Polaris to optimise any differentiable objective function.

The event handler mechanism (`Axon.Loop.handle/4`) lets you attach callbacks to loop events such as iteration completion or epoch end. The README mentions plans for distributed training loops in future releases but provides no implementation timeline.

ONNX Integration and Platform Limits

Axon does not natively run ONNX models. The README points to two external projects for ONNX interoperability: Ortex, which runs ONNX models directly via ONNX Runtime bindings, and AxonONNX, which converts ONNX models to Axon models where conversion is possible.

The distinction matters: Ortex runs ONNX models as-is, preserving any operators the ONNX Runtime supports. AxonONNX converts the ONNX model into Axon's graph representation, which enables full integration with Nx backends but only works when every ONNX operator has an Axon equivalent.

For users who want to export Axon models, the README notes that because a model is just an Elixir struct, serialising it to multiple formats is straightforward in principle, though the README does not document specific export targets beyond ONNX conversion.

Comparison With PyTorch and TensorFlow

PyTorch is the dominant framework for model training in Python. Its training loop philosophy, which Axon's training API is described as inspired by through PyTorch Ignite, involves manually writing loops or using high-level trainers like those in PyTorch Lightning. Both Axon and PyTorch allow you to JIT-compile models, though PyTorch does this with `torch.jit` and Axon does it through Nx's compiler backend system.

TensorFlow Keras is the other reference point in the README: the model creation API is described as targeting parity with the layer types available in frameworks like PyTorch or TensorFlow Keras. The difference in Axon is that layers are Elixir functions rather than Python objects, and the training state is managed separately from the model definition.

The practical implication is that neither PyTorch nor TensorFlow models can run in Axon without conversion through AxonONNX or Ortex. Teams with existing Python model code cannot run it in Axon without this conversion step.

Maintenance, Releases, and License

The repository is not archived. The last push was on 2026-09-21, which coincides with the v0.9.0 release on the same date. Prior releases include v0.8.1 on 2026-03-11 and v0.8.0 on 2025-11-14, indicating roughly quarterly releases.

The project is licensed under Apache-2.0. Dockyard is listed in the README as a sponsor. Axon was created by Sean Moriarity, as stated in the copyright notice in the LICENSE file.

The README notes that support for distributed training loops is planned and that the project is also seeking ways to improve training loop performance by running entirely on native accelerators, though neither feature is implemented as of the documented state.

Editorial conclusion

Axon is the natural choice for teams building neural networks in Elixir who are already using the Nx ecosystem. Teams with large PyTorch or TensorFlow codebases will find the bridge cost high. Developers who want ONNX interoperability should evaluate Ortex or AxonONNX alongside Axon. Verify that the Nx backend you plan to use is supported before committing, as EXLA is the primary accelerated backend documented in the README.

Frequently asked questions

Official sources

1. elixir-nx/axon on GitHub
2. Issues
3. License: Apache-2.0
4. README
5. Releases