dantinox — Top-Level API

The top-level dantinox package exposes a three-level API. Import dantinox as dx to access all public symbols.

---

Low-code functional API

These functions are the Level-1 entry points — no class instantiation required.

dantinox.fit(paradigm: str, data_source: str, *, run_dir: str | None = None, training_config: TrainingConfig | None = None, **kwargs: Any) → str

One-call shortcut: build paradigm, train, return run directory.

Keyword arguments that match ModelConfig fields are forwarded to the config constructor; everything else goes to TrainingConfig.

Example:

run_dir = dx.fit("ar", "data/wiki.txt",
                 dim=512, n_heads=8, head_size=64, num_blocks=12,
                 vocab_size=32_000, lr=3e-4, epochs=5)

dantinox.train(paradigm: Paradigm, data_source: str | None = None, *, run_dir: str | None = None, training_config: TrainingConfig | None = None, **training_kwargs: Any) → str

Train paradigm on data_source and return the run directory.

Parameters:

• paradigm – Any Paradigm instance.

• data_source – Path to a text file, or a HuggingFace dataset name. May be omitted when training_config sets dataset_source="huggingface" and dataset_name.

• run_dir – Output directory (auto-generated when omitted).

• training_config – A ready-made TrainingConfig. When supplied, training_kwargs are ignored.

• **training_kwargs – Forwarded to TrainingConfig — e.g. lr=3e-4, epochs=10, batch_size=64.

Returns:

Absolute path to the run directory with the best checkpoint.

Example:

run_dir = dx.train(paradigm, "data/wiki.txt", lr=1e-4, epochs=3)

cfg = dx.TrainingConfig(lr=1e-4, epochs=3, dataset_source="huggingface",
                        dataset_name="wikitext")
run_dir = dx.train(paradigm, training_config=cfg)

dantinox.build(paradigm: str, config: ModelConfig | None = None, **model_kwargs: Any) → Paradigm

Construct a Paradigm from a string name and optional ModelConfig.

This is a thin wrapper around Paradigm(ModelConfig(paradigm=..., ...)). For new code, prefer the explicit form:

p = dx.Paradigm(dx.ModelConfig(paradigm="ar", dim=512, n_heads=8,
                                num_blocks=12, vocab_size=32_000))

Parameters:

• paradigm – "ar" | "discrete" | "continuous" | "embedder"

• config – A ModelConfig. When omitted, model_kwargs are forwarded to ModelConfig.

• **model_kwargs – Forwarded to ModelConfig when config is None.

Returns:

A ready-to-use Paradigm instance.

Example:

p = dx.build("ar", dim=512, n_heads=8, num_blocks=12, vocab_size=32_000)

p = dx.build("continuous", dim=256, n_heads=4, num_blocks=4,
             embed_dim=768, bottleneck_dim=128)

dantinox.profile(config: ModelConfig, seq_len: int, batch_size: int = 1, *, n_warmup: int = 5, n_runs: int = 20, model: Any = None) → ProfilingResult

Profile a model: FLOPs + latency + throughput.

When model is provided, a real JAX forward pass is timed. Otherwise only the analytical FLOPs estimate is returned (latency = 0).

Example:

cfg    = dx.ModelConfig(dim=512, n_heads=8, head_size=64, num_blocks=12,
                        vocab_size=32_000)
report = dx.profile(cfg, seq_len=512, batch_size=4)
print(report.flops)
print(report.latency)

dantinox.load(run_dir: str, paradigm: Paradigm | None = None, checkpoint: str = 'best') → Any

Load a checkpoint from run_dir and return the NNX model.

Parameters:

• run_dir – Directory produced by Trainer.fit() or dx.fit().

• paradigm – Optional Paradigm instance. When omitted, config.yaml in run_dir is read to reconstruct the ModelConfig and infer the paradigm automatically. Pass explicitly to override or for legacy checkpoints without a config.yaml.

• checkpoint –

  Which checkpoint to load. One of:

  • "best" (default) — checkpoint_best.msgpack

  • "latest" — checkpoint_latest.msgpack

  • any other string — checkpoint_<string>.msgpack

  • an absolute or relative path ending in .msgpack to load a file directly.

Example:

model = dx.load("runs/20240101_120000")
model = dx.load("runs/20240101_120000", checkpoint="latest")
model = dx.load("runs/20240101_120000", checkpoint="/path/to/my.msgpack")

dantinox.quick_generate(run_dir: str, prompt: str, *, paradigm: Paradigm | None = None, tokenizer: Any = None, max_new_tokens: int = 200, temperature: float = 1.0, use_cache: bool = True) → str

Load checkpoint and generate text — no boilerplate required.

When paradigm is given the checkpoint is restored through paradigm.build_model() and decoded with paradigm.generate; tokenizer (or the run’s saved tokenizer.json) handles text ↔ ids. Otherwise the run directory is loaded with Generator.

Example:

print(dx.quick_generate("runs/20240101_120000", "Once upon a time"))

---

Re-exported symbols

The following symbols are importable directly from dantinox:

Configs

• ModelConfig — model architecture configuration

• TrainingConfig — training hyperparameters

• Config — legacy unified config (backward-compat)

ELFConfig (flow-matching architecture config, aliased to FlowMatchingConfig) is not re-exported at the top level — import it from the submodule instead: from dantinox.core.config import FlowMatchingConfig (or the deprecated alias ELFConfig from the same module).

Paradigms

• Paradigm — unified paradigm dispatcher (selects AR/discrete/continuous/embedder from ModelConfig.paradigm)

• ARParadigm — autoregressive

• DiscreteParadigm — masked diffusion (LLaDA formulation)

• ContinuousParadigm — continuous flow-matching

• EmbedderParadigm — contrastive text-embedding paradigm

DiscreteConfig (masked-diffusion hyperparameters) is not re-exported at the top level — import it directly: from dantinox.paradigms.diffusion.discrete import DiscreteConfig.

Training

• Trainer — paradigm-agnostic training harness

• build_optimizer — optimizer factory

• build_schedule — LR schedule factory

Profiling

• count_flops — analytical FLOPs estimator

• FLOPsBreakdown — per-component FLOPs result

• LatencyTracker — wall-clock latency measurement

• ProfilingResult — aggregated latency statistics

• profile_fn — functional latency wrapper

Benchmarking

• BenchmarkSuite — task orchestrator

• BenchmarkTask — plugin base class

• BenchmarkConfig — suite configuration

• BenchmarkResult — per-task result

• SuiteReport — aggregated report

• ThroughputTask, LatencyTask, PerplexityTask — built-in tasks

Visualization

• Visualizer — chart registry and renderer

• Chart — chart ABC

• RenderConfig — rendering options

• TrainingCurveChart, ThroughputChart, ThroughputBatchChart — built-in charts

• LatencyChart, RadarChart, ParetoChart — built-in charts