DantinoX

A research-grade JAX/Flax NNX library for language model research. Three generation paradigms — Autoregressive, Masked Diffusion, and Continuous Flow-Matching — 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. Discrete Diffusion vs. Continuous Flow-Matching.
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.
Discrete Diffusion: 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. Follows the LLaDA formulation. See Discrete Diffusion.
Continuous Flow-Matching: Operates in the continuous embedding space. Transforms Gaussian noise into clean token embeddings via an Euler ODE solver. See Continuous Flow-Matching.

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 and Fontana, Aleksandar and Rossolini, Giulio
             and Saracino, Andrea},
  title   = {{D}antino{X}: A Unified Framework for Multi-Paradigm
             Language Modeling},
  year    = {2026},
  url     = {https://github.com/winstonsmith1897/DantinoX},
}