Overview

LANfactory

lanfactory is a lightweight Python package for training likelihood approximation networks (LANs) — and their choice-probability siblings — for sequential sampling models (SSMs), using PyTorch or JAX/Flax. Starting from simulator-generated training data, it provides dataloaders, network factories, and training loops, and exports the trained networks to ONNX so they can serve as likelihoods in HSSM.

New to the idea? Read likelihood approximation networks and LANfactory before choosing a training backend.

---

Installation

pip install lanfactory

Optional integrations ship as extras: lanfactory[mlflow] (experiment tracking), lanfactory[hf] (HuggingFace Hub upload/download), lanfactory[sbi] and lanfactory[bayesflow] (ONNX export of externally trained networks), or lanfactory[all] for everything.

---

Ecosystem fit

LANfactory is the network-training layer of the HSSM ecosystem: it trains LAN, CPN, OPN, and gonogo networks on simulated data and exports them to ONNX in the form that HSSM consumes as likelihoods.

For the full map — what each package owns, how artifacts flow between them, and which versions work together — see The HSSM ecosystem.

---

Quickstart

Given a folder of training data files generated with ssm-simulators, the minimal PyTorch training loop is shown below. This abbreviated API sketch assumes those files already exist; the linked tutorial is the first-success path from data generation through a trained network.

from pathlib import Path
import lanfactory

file_list = list(Path("training_data").glob("*.pickle"))

train_dl, valid_dl, input_dim = lanfactory.trainers.make_train_valid_dataloaders(
    file_ids=file_list, batch_size=128, network_type="lan"
)
net = lanfactory.trainers.TorchMLPFactory(
    network_config=lanfactory.config.network_configs.network_config_mlp,
    input_dim=input_dim,
    network_type="lan",
)
trainer = lanfactory.trainers.ModelTrainerTorchMLP(
    model=net,
    train_config=lanfactory.config.network_configs.train_config_mlp,
    train_dl=train_dl,
    valid_dl=valid_dl,
)
trainer.train_and_evaluate(output_folder="torch_models/ddm", output_file_id="ddm")

For the full walkthrough — data generation, configuration, training, and inspecting the learned likelihood — see the training tutorial.

---

Export to ONNX

Trained PyTorch networks convert to ONNX with the transform-onnx CLI:

transform-onnx --network-config-file <network_config.pickle> \
  --state-dict-file <state_dict.pt> --input-shape <input_dim> \
  --output-onnx-file <model.onnx>

The resulting file can be used directly with HSSM — see The ONNX likelihood contract for the artifact rules. Networks trained outside LANfactory can be exported the same way — see the sbi and bayesflow export guides.

---

Where to go next

• Tutorials
  • Train a network (PyTorch LAN) — the canonical end-to-end training walkthrough.
  • How to train with the JAX backend — the same workflow on JAX/Flax.
  • Exporting sbi → ONNX and exporting bayesflow → ONNX — runnable export notebooks.
• Guides
  • Network types: LAN, CPN, OPN — what each network learns and how their configs differ.
  • MLflow integration — track and compare training runs.
  • HuggingFace Hub — upload and download trained networks.
  • Exporting sbi models and exporting bayesflow models — bring externally trained networks into HSSM.
• API reference — config, trainers, ONNX, network inspectors, Hugging Face, and utilities.

We hope this package may be helpful in case you attempt to train LANs for your own research.