Axon: Nx-Powered Neural Networks in Elixir

The gap Axon fills between Nx and a training loop

Nx gives Elixir tensors, numerical definitions and compilers. It does not give you layers, initializers, loss functions or a training loop. Axon is the layer that sits on top. The README lists five functional modules (Axon.Activations, Axon.Initializers, Axon.Layers, Axon.Losses, Axon.Metrics) plus a model creation API and a training API. The intended reader is an Elixir engineer who wants to define a network in the same language as the rest of the application, rather than exporting weights to a Python process. The README's stated goal is an API that is productive, extensible and on par with other modern deep learning frameworks, and it invites issues for missing functionality, which is an admission that the surface is not complete. That invitation is the honest framing: Axon is a framework for people who have already chosen Elixir and want the model to live there too.

Three APIs that do not depend on each other

The separation is the design decision worth understanding. The functional API is a set of modules implemented as numerical definitions, which the README says means you can use any Nx compiler or backend to accelerate them, and you can compose them with your own defn functions. The model creation API manages initialization and application only. The README states the philosophy directly: a model's only concerns are initialization and application, and the model should not be concerned with training. Axon.build/2 turns the struct into two functions, an init_fn and a predict_fn. The training API is a separate layer inspired by PyTorch Ignite. Because the layers are numerical definitions rather than opaque kernels, the same network can be JIT or AOT compiled, or transformed into formats like TensorFlow Lite and ONNX, per the README. The cost of this decoupling is that you own more wiring: Axon.build/2 returns functions you call yourself, and if you skip Axon.Loop you write the step function.

What a model actually is: an Elixir struct

A model is built by piping Axon.input/2 into layer functions. The README's example is Axon.input("input", shape: {nil, 784}) piped into Axon.dense(128) and then Axon.dense(10, activation: :softmax). The shape tuple uses nil for the batch dimension, so the model is batch-size agnostic at definition time. The result is an Elixir struct, and the README notes that serializing it to multiple formats in the future is straightforward because of that. Inspection is built in: the default inspect protocol gives a summary, and Axon.Display.as_table/2 prints a table with columns for Layer, Input Shape, Output Shape, Options and Parameters. The README's rendered table shows the softmax step appearing as its own row, softmax_0 ( softmax["dense_1"] ), rather than being folded into the dense layer, and it reports Total Parameters and Total Parameters Memory at the bottom. For a two-layer dense network over 784 inputs, the table lists 101770 parameters. That number is a useful sanity check when you are assembling a model by hand and want to confirm the shapes line up before you spend time on a training run.

Getting it running: build, initialize, predict

The README gives the full sequence. Define the model, then destructure the build output. {init_fn, predict_fn} = Axon.build(model, compiler: EXLA) is the call, with the compiler passed as an option. Initialization takes a template and an initial parameter map: params = init_fn.(Nx.template({1, 784}, :f32), %{}). Prediction then takes the parameters and the input: predict_fn.(params, input). Note that Nx.template is given a concrete batch size of 1 for initialization even though the model was defined with a nil batch dimension. For training, the README's pattern is four steps: define the model, define a loop with a factory method, instrument it with metrics and handlers, run it on data. The factory call in the example is Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(0.005)), followed by Axon.Loop.metric(:accuracy), then Axon.Loop.handle(:iteration_completed, &log_metrics/1, every: 50), then Axon.Loop.run(data, %{}, epochs: 10, compiler: EXLA). The run returns an %Axon.Loop.State{} whose step_state holds model_state, which is how you retrieve trained parameters. The optimizer comes from Polaris, not Axon, and the README stresses that the optimization API does not depend on Axon models and can optimize any differentiable objective.

Where Axon is the wrong choice

The README ends mid-sentence on a future plan to support distribution, so as of the documented state, multi-node training is not a described capability. If your workload needs to span machines, this is not the tool yet. The training API is also a convenience layer, not a production serving system: the documented flow runs a loop over data and returns state, and there is no described checkpointing, metrics export or model server in the material. Layer coverage is another boundary. The README says Axon has support for the same high-level layers you would find in PyTorch or TensorFlow Keras, but it also asks users to open an issue if functionality is missing, which means the set is not guaranteed complete. When a layer is absent from the high-level API, the fallback is Axon.Layers or your own defn, and that fallback requires you to reason about shapes and parameter names yourself. Finally, the model-struct design means training concerns are deliberately out of scope, so anything you want around the loop is code you write.

How this differs from doing it in Python

The obvious alternative is PyTorch or TensorFlow in Python, and the difference is not the layer list, which the README claims is comparable. The difference is where the model lives. With Axon, the model is an Elixir struct and the layers are defn functions, so the network composes with the rest of an Elixir system and can be compiled by an Nx backend or exported to TensorFlow Lite and ONNX. With Python frameworks, the model lives in a Python process and the Elixir side talks to it over a boundary. The trade is ecosystem depth. Python frameworks have years of pretrained weights, dataset tooling and deployment targets that Axon does not describe. If your problem is standard and you want the shortest path to a trained model, Python is still the safer bet. Axon is for the case where the model is a component inside an Elixir application and the language boundary is the thing you want to remove.

Licence and the Polaris dependency

Axon is Apache-2.0, which permits commercial use and modification under the terms of that licence; this is a description of the identifier, not legal advice, and you should read the licence text for your own situation. The dependency worth tracking is Polaris, which supplies the optimization API. The README is explicit that Axon uses Polaris for optimization and that the optimization API does not depend on Axon models. That means a version bump in Polaris can affect your training loop even when Axon itself is unchanged, and it means you can use Polaris for objectives that have nothing to do with Axon. The release history in the repository metadata shows v0.7.0 in October 2024, v0.8.0 in November 2025 and v0.8.1 in March 2026, so the cadence between the 0.7 and 0.8 lines was over a year. Plan for a library that moves in larger steps rather than continuous patches, and pin both Axon and Polaris in mix.lock when you take a version.

Editorial conclusion

Adopt Axon if your stack is already Elixir and you want to define models as structs, compile them with an Nx backend such as EXLA, and drive training through Axon.Loop without leaving the language. Do not adopt it if you need distribution across nodes or a serving runtime, since the README only lists distribution as a future plan. Before committing, verify which Nx compiler or backend you will target, confirm that the layer you need exists in Axon.Layers rather than the high-level API, and check that pinning Polaris alongside Axon does not conflict in your mix.lock.

Official sources

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