# `Edifice`
[🔗](https://github.com/blasphemetheus/edifice/blob/main/lib/edifice.ex#L1)

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

| Family | Architectures |
|--------|--------------|
| Transformer | Decoder-Only (GPT-style), Multi-Token Prediction, Byte Latent Transformer, Nemotron-H |
| Feedforward | MLP, KAN, KAT, TabNet, BitNet |
| Convolutional | Conv1D/2D, ResNet, DenseNet, TCN, MobileNet, EfficientNet |
| Recurrent | LSTM, GRU, xLSTM, xLSTM v2, mLSTM, sLSTM, MinGRU, MinLSTM, DeltaNet, Gated DeltaNet, TTT, TTT-E2E, Titans, Reservoir (ESN), Native Recurrence |
| State Space | Mamba, Mamba-2 (SSD), Mamba-3, S4, S4D, S5, H3, Hyena, Hyena v2, BiMamba, GatedSSM, GSS, StripedHyena, Hymba, State Space Transformer |
| Attention | Multi-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 |
| Vision | ViT, DeiT, Swin, U-Net, ConvNeXt, MLP-Mixer, FocalNet, PoolFormer, NeRF, MambaVision |
| Generative | VAE, VQ-VAE, GAN, Diffusion, DDIM, DiT, DiT v2, MMDiT, Latent Diffusion, Consistency, Score SDE, Flow Matching, SoFlow, Normalizing Flow, Transfusion, CogVideoX, TRELLIS |
| Graph | GCN, GAT, GraphSAGE, GIN, GINv2, PNA, GraphTransformer, SchNet, Message Passing |
| Sets | DeepSets, PointNet |
| Energy | EBM, Hopfield, Neural ODE |
| Probabilistic | Bayesian, MC Dropout, Evidential |
| Memory | NTM, Memory Networks |
| Meta | MoE, 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 |
| Liquid | Liquid Neural Networks |
| Contrastive | SimCLR, BYOL, Barlow Twins, MAE, VICReg, JEPA, Temporal JEPA |
| Interpretability | Sparse Autoencoder, Transcoder |
| World Model | World Model |
| Multimodal | MLP Projection Fusion, Cross-Attention Fusion, Perceiver Resampler |
| RL | PolicyValue |
| Neuromorphic | SNN, ANN2SNN |
| Inference | Medusa |
| Robotics | ACT, OpenVLA |
| Audio | SoundStorm, EnCodec, VALL-E |

# `build`

```elixir
@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`

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

List all available architecture names.

## Examples

    iex> :mamba in Edifice.list_architectures()
    true

# `list_families`

```elixir
@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`

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

Get the module for a named architecture.

## Examples

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

---

*Consult [api-reference.md](api-reference.md) for complete listing*