DantinoX

A research-grade JAX/Flax NNX library for language model research. Three generation paradigms — Autoregressive, Masked Diffusion, and ELF — on the same Transformer architecture, with a single trainer and zero boilerplate.

Overview

DantinoX was created to answer a single question: how do different generation paradigms — autoregressive, masked diffusion, and flow-matching — compare when trained on the same architecture with the same training code?

The library targets three audiences:

Researchers who want a reproducible comparison of AR vs. Diffusion vs. ELF.
Students who want to read the internals of a modern Transformer.
Engineers who need architectural variants (GQA, MLA, MoE, LoRA) without rewriting the trainer.

Three Generation Paradigms

Autoregressive (AR): The classical left-to-right paradigm. Generates one token at a time using a causal (masked) attention pattern and a static pre-allocated KV-cache. See Autoregressive (AR) Generation.
Masked Diffusion (LLaDA): Generates all tokens in parallel from a fully masked sequence and iteratively unmasks them over multiple diffusion steps. Attention is bidirectional. Optionally accelerated by Fast-dLLM DualCache. See Discrete Diffusion.
ELF — Continuous Flow Matching: Operates in the continuous embedding space. Transforms Gaussian noise into clean token embeddings via an Euler ODE solver. See paradigms/elf.

Quick Install

pip install dantinox          # CPU / GPU (JAX auto-detected)

Or from source:

git clone https://github.com/winstonsmith1897/DantinoX
cd DantinoX && pip install -e ".[dev]"

One-Liner Usage

import dantinox as dx

run_dir = dx.fit("ar", "data/wiki.txt",
                 dim=512, n_heads=8, head_size=64,
                 num_blocks=12)

print(dx.quick_generate(run_dir, "In the beginning"))

Citation

@software{dantinox2026,
  author  = {Simoni, Marco},
  title   = {DantinoX: A Unified {JAX}/Flax Framework for {AR},
             Masked Diffusion, and Flow-Matching Language Models},
  year    = {2026},
  url     = {https://github.com/winstonsmith1897/DantinoX},
}