Edifice (Edifice v0.2.0)

Copy Markdown View Source

Edifice - A comprehensive ML architecture library for Elixir.

Provides implementations of all major neural network architecture families, built on Nx and Axon. From simple MLPs to state space models, attention mechanisms, generative models, and graph neural networks.

Quick Start

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

# Or use the module directly
model = Edifice.SSM.Mamba.build(embed_dim: 256, hidden_size: 512)

# List all available architectures
Edifice.list_architectures()

Architecture Families

FamilyArchitectures
TransformerDecoder-Only (GPT-style), Multi-Token Prediction, Byte Latent Transformer, Nemotron-H
FeedforwardMLP, KAN, KAT, TabNet, BitNet
ConvolutionalConv1D/2D, ResNet, DenseNet, TCN, MobileNet, EfficientNet
RecurrentLSTM, GRU, xLSTM, xLSTM v2, mLSTM, sLSTM, MinGRU, MinLSTM, DeltaNet, Gated DeltaNet, TTT, TTT-E2E, Titans, Reservoir (ESN), Native Recurrence
State SpaceMamba, Mamba-2 (SSD), Mamba-3, S4, S4D, S5, H3, Hyena, Hyena v2, BiMamba, GatedSSM, GSS, StripedHyena, Hymba, State Space Transformer
AttentionMulti-Head, GQA, MLA, KDA (Kimi Delta Attention), DiffTransformer, Perceiver, FNet, Linear Transformer, Nystromformer, Performer, RetNet, RetNet v2, RWKV, GLA, GLA v2, HGRN, HGRN v2, Griffin, Hawk, Based, InfiniAttention, Conformer, Mega, MEGALODON, RingAttention, Lightning Attention, Flash Linear Attention, YaRN, NSA, Dual Chunk Attention
VisionViT, DeiT, Swin, U-Net, ConvNeXt, MLP-Mixer, FocalNet, PoolFormer, NeRF, MambaVision
GenerativeVAE, VQ-VAE, GAN, Diffusion, DDIM, DiT, DiT v2, MMDiT, Latent Diffusion, Consistency, Score SDE, Flow Matching, SoFlow, Normalizing Flow, Transfusion, CogVideoX, TRELLIS
GraphGCN, GAT, GraphSAGE, GIN, GINv2, PNA, GraphTransformer, SchNet, Message Passing
SetsDeepSets, PointNet
EnergyEBM, Hopfield, Neural ODE
ProbabilisticBayesian, MC Dropout, Evidential
MemoryNTM, Memory Networks
MetaMoE, MoE v2, Switch MoE, Soft MoE, LoRA, DoRA, Adapter, Hypernetworks, Capsules, MixtureOfDepths, MixtureOfAgents, RLHFHead, Speculative Decoding, Test-Time Compute, Mixture of Tokenizers, Speculative Head, Distillation Head, QAT, Hybrid Builder
LiquidLiquid Neural Networks
ContrastiveSimCLR, BYOL, Barlow Twins, MAE, VICReg, JEPA, Temporal JEPA
InterpretabilitySparse Autoencoder, Transcoder
World ModelWorld Model
MultimodalMLP Projection Fusion, Cross-Attention Fusion, Perceiver Resampler
RLPolicyValue
NeuromorphicSNN, ANN2SNN
InferenceMedusa
RoboticsACT, OpenVLA
AudioSoundStorm, EnCodec, VALL-E

Summary

Functions

build(name, opts \\ [])

Build an architecture by name.

list_architectures()

List all available architecture names.

list_families()

List architectures grouped by family.

module_for(name)

Get the module for a named architecture.

Functions

build(name, opts \\ [])

@spec build(
  atom(),
  keyword()
) :: Axon.t() | tuple()

Build an architecture by name.

Parameters

  • name - Architecture name (see list_architectures/0)

  • opts - Architecture-specific options. The primary input dimension option varies by family:

    • :embed_dim — sequence models (SSM, attention, recurrent)
    • :input_size — flat-vector models (MLP, probabilistic, energy)
    • :input_dim — graph/spatial models (GCN, DeepSets, vision)
    • :obs_size — diffusion models (Diffusion, DDIM, FlowMatching)

    You can pass any of the first three interchangeably — build/2 will normalize to the name each module expects. :obs_size is also normalized from the above. Image models requiring :input_shape (a tuple like {nil, 32, 32, 3}) must be passed explicitly.

Examples

model = Edifice.build(:mamba, embed_dim: 256, hidden_size: 512, num_layers: 4)
model = Edifice.build(:mlp, input_size: 256, hidden_sizes: [512, 256])
model = Edifice.build(:gcn, input_dim: 8, hidden_dims: [32, 32], num_classes: 10)

Returns

An Axon.t() model for most architectures. These architectures return tuples:

  • :vae — {encoder, decoder}
  • :vq_vae — {encoder, decoder}
  • :gan — {generator, discriminator}
  • :normalizing_flow — {flow_model, log_det_fn}
  • :simclr — {backbone, projection_head}
  • :byol — {online_network, target_network}
  • :barlow_twins — {backbone, projection_head}
  • :vicreg — {backbone, projection_head}
  • :jepa — {context_encoder, predictor}
  • :temporal_jepa — {context_encoder, predictor}
  • :mae — {encoder, decoder}
  • :world_model — {encoder, dynamics, reward_head} (or 4-tuple with decoder)
  • :byte_latent_transformer — {encoder, latent_transformer, decoder}
  • :speculative_decoding — {draft_model, verifier_model}
  • :multi_token_prediction — Axon.container(%{pred_1: ..., pred_N: ...})
  • :test_time_compute — Axon.container(%{backbone: ..., scores: ...})
  • :speculative_head — Axon.container(%{pred_1: ..., pred_N: ...})
  • :act — {encoder, decoder}
  • :medusa — Axon.container(%{head_1: ..., head_K: ...})

list_architectures()

@spec list_architectures() :: [atom()]

List all available architecture names.

Examples

iex> :mamba in Edifice.list_architectures()
true

list_families()

@spec list_families() :: %{required(atom()) => [atom()]}

List architectures grouped by family.

Examples

Edifice.list_families()
# => %{
#   feedforward: [:mlp, :kan],
#   ssm: [:mamba, :mamba_ssd, :s5, ...],
#   ...
# }

module_for(name)

@spec module_for(atom()) :: module()

Get the module for a named architecture.

Examples

Edifice.module_for(:mamba)
# => Edifice.SSM.Mamba