Edifice

A comprehensive ML architecture library for Elixir, built on Nx and Axon.

92 neural network architectures across 16 families — from MLPs to Mamba, transformers to graph networks, VAEs to spiking neurons.

Why Edifice?

The Elixir ML ecosystem has excellent numerical computing (Nx) and model building (Axon) foundations, but no comprehensive collection of ready-to-use architectures. Edifice fills that gap:

• One dependency for all major architecture families
• Consistent API — every architecture follows Module.build(opts) returning an Axon model
• Unified registry — Edifice.build(:mamba, opts) discovers and builds any architecture by name
• Pure Elixir — no Python, no ONNX imports, just Nx/Axon all the way down
• GPU-ready — works with EXLA/CUDA out of the box

Installation

Add edifice to your dependencies in mix.exs:

def deps do
  [
    {:edifice, "~> 0.1.0"}
  ]
end

Edifice requires Nx ~> 0.10 and Axon ~> 0.8. For GPU acceleration, add EXLA:

{:exla, "~> 0.10"}

Quick Start

# Build any architecture by name
model = Edifice.build(:mamba, embed_size: 256, hidden_size: 512, num_layers: 4)

# Or use the module directly for more control
model = Edifice.SSM.Mamba.build(
  embed_size: 256,
  hidden_size: 512,
  state_size: 16,
  num_layers: 4,
  window_size: 60
)

# Build and run
{init_fn, predict_fn} = Axon.build(model)
params = init_fn.(Nx.template({1, 60, 256}, :f32), Axon.ModelState.empty())
output = predict_fn.(params, input)

# Explore what's available
Edifice.list_architectures()
# => [:attention, :bayesian, :capsule, :deep_sets, :densenet, :diffusion, ...]

Edifice.list_families()
# => %{ssm: [:mamba, :mamba_ssd, :s5, ...], attention: [:attention, :retnet, ...], ...}

Architecture Families

Feedforward

State Space Models

Attention & Linear Attention

Recurrent Networks

Vision

Convolutional

Generative Models

Contrastive & Self-Supervised

Graph & Set Networks

Energy, Probabilistic & Memory

Meta-Learning & Specialized

Building Blocks

Guides

New to ML?

Start here if you're new to machine learning. These guides build from zero to fluency with Edifice's API and architecture families.

1. ML Foundations — What neural networks are, how they learn, tensors and shapes
2. Core Vocabulary — Essential terminology used across all guides
3. The Problem Landscape — Classification, generation, sequence modeling — which architectures solve which problems
4. Reading Edifice — The build/init/predict pattern, Axon graphs, shapes, and runnable examples
5. Learning Path — A guided tour through the 19 architecture families

Architecture Guides

Conceptual guides covering theory, architecture evolution, and decision tables for each family.

Sequence Processing

• State Space Models — S4 through Mamba to hybrid architectures
• Attention Mechanisms — Quadratic to linear to Fourier to retention
• Recurrent Networks — LSTM through xLSTM, MinGRU, TTT, and Titans

Representation Learning

• Vision Architectures — ViT, Swin, UNet, ConvNeXt, MLP-Mixer
• Convolutional Networks — ResNet, DenseNet, MobileNet, TCN
• Contrastive Learning — SimCLR, BYOL, BarlowTwins, MAE, VICReg
• Graph & Set Networks — Message passing, spectral, invariance

Generative & Dynamic

• Generative Models — VAEs, GANs, diffusion, flows
• Dynamic & Continuous — ODE dynamics, energy landscapes, spiking

Composition & Enhancement

• Building Blocks — RoPE vs ALiBi, RMSNorm, SwiGLU, composition
• Meta-Learning — MoE, PEFT (LoRA/Adapter), capsules
• Uncertainty & Memory — Bayesian, NTM, MLP/KAN/TabNet foundations

Examples

See examples/ for runnable scripts including mlp_basics.exs, sequence_comparison.exs, graph_classification.exs, vae_generation.exs, and architecture_tour.exs.

Mamba for Sequence Modeling

model = Edifice.SSM.Mamba.build(
  embed_size: 128,
  hidden_size: 256,
  state_size: 16,
  num_layers: 4,
  window_size: 100
)

{init_fn, predict_fn} = Axon.build(model)
params = init_fn.(Nx.template({1, 100, 128}, :f32), Axon.ModelState.empty())
output = predict_fn.(params, Nx.broadcast(0.5, {1, 100, 128}))
# => {1, 256}

Graph Classification with GCN

model = Edifice.Graph.GCN.build_classifier(
  input_dim: 16,
  hidden_dims: [64, 64],
  num_classes: 2,
  pool: :mean
)

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

params = init_fn.(
  %{
    "nodes" => Nx.template({4, 10, 16}, :f32),
    "adjacency" => Nx.template({4, 10, 10}, :f32)
  },
  Axon.ModelState.empty()
)

output = predict_fn.(params, %{
  "nodes" => Nx.broadcast(0.5, {4, 10, 16}),
  "adjacency" => Nx.eye(10) |> Nx.broadcast({4, 10, 10})
})
# => {4, 2}

VAE with Reparameterization

{encoder, decoder} = Edifice.Generative.VAE.build(
  input_size: 784,
  latent_size: 32,
  encoder_sizes: [512, 256],
  decoder_sizes: [256, 512]
)

# Encoder outputs mu and log_var
{init_fn, predict_fn} = Axon.build(encoder)
params = init_fn.(Nx.template({1, 784}, :f32), Axon.ModelState.empty())
%{mu: mu, log_var: log_var} = predict_fn.(params, Nx.broadcast(0.5, {1, 784}))

# Sample latent vector (requires PRNG key for stochastic sampling)
key = Nx.Random.key(42)
{z, _new_key} = Edifice.Generative.VAE.reparameterize(mu, log_var, key)

# KL divergence for training
kl_loss = Edifice.Generative.VAE.kl_divergence(mu, log_var)

Permutation-Invariant Set Processing

model = Edifice.Sets.DeepSets.build(
  input_dim: 3,
  hidden_dim: 64,
  output_dim: 10,
  pool: :mean
)

{init_fn, predict_fn} = Axon.build(model)
params = init_fn.(Nx.template({4, 20, 3}, :f32), Axon.ModelState.empty())
# Process sets of 20 3D points
output = predict_fn.(params, Nx.broadcast(0.5, {4, 20, 3}))
# => {4, 10}

API Design

Every architecture module follows the same pattern:

# Module.build(opts) returns an Axon model
model = Edifice.SSM.Mamba.build(embed_size: 256, hidden_size: 512)

# Some modules expose layer-level builders for composition
layer = Edifice.Graph.GCN.gcn_layer(nodes, adjacency, output_dim)

# Generative models may return tuples
{encoder, decoder} = Edifice.Generative.VAE.build(input_size: 784)

# Utility functions for training
loss = Edifice.Generative.VAE.loss(reconstruction, target, mu, log_var)
energy = Edifice.Energy.Hopfield.energy(query, patterns, beta)

The unified registry lets you build any architecture by name:

# Useful for hyperparameter search, config-driven experiments
for arch <- [:mamba, :retnet, :griffin, :gla] do
  model = Edifice.build(arch, embed_size: 256, hidden_size: 512, num_layers: 4)
  # ... train and evaluate
end

Requirements

• Elixir >= 1.18
• Nx ~> 0.10
• Axon ~> 0.8
• Polaris ~> 0.1
• EXLA ~> 0.10 (optional, for GPU acceleration)

License

MIT License. See LICENSE for details.