trainers

lanfactory.trainers

Modules:

• jax_mlp –

  This module contains the JaxMLP class and the ModelTrainerJaxMLP class which

• torch_mlp –

  This module contains the classes for training TorchMLP models.

Classes:

• DatasetTorch –

  Dataset class for TorchMLP training.

• JaxMLP –

  JaxMLP class.

• LoadTorchMLP –

  General-purpose class to load TorchMLP models.

• LoadTorchMLPInfer –

  Model loader with inference mode enabled by default.

• ModelTrainerJaxMLP –

  Class for training JaxMLP models.

• ModelTrainerTorchMLP –
• TorchMLP –

  TorchMLP class.

Functions:

• JaxMLPFactory –

  Factory function to create a JaxMLP object.

• TorchMLPFactory –

  Factory function to create a TorchMLP object.

• make_dataloader –

  Create a DataLoader for LAN/CPN/OPN training.

• make_train_valid_dataloaders –

  Create train and validation DataLoaders with automatic file splitting.

lanfactory.trainers.DatasetTorch

DatasetTorch(file_ids: list[str] | list[Path], batch_size: int = 32, label_lower_bound: float | None = None, label_upper_bound: float | None = None, features_key: str = 'data', label_key: str = 'labels', out_framework: str = 'torch')

Bases: Dataset

Dataset class for TorchMLP training.

Arguments

file_ids (list[str]):
    List of paths to the data files.
batch_size (int):
    Batch size.
label_lower_bound (float | None):
    Lower bound for the labels.
label_upper_bound (float | None):
    Upper bound for the labels.
features_key (str):
    Key for the features in the data files.
label_key (str):
    Key for the labels in the data files.
out_framework (str):
    Output framework.

lanfactory.trainers.JaxMLP

Bases: Module

JaxMLP class.

Arguments

layer_sizes (Sequence[int]):
    Sequence of integers containing the sizes of the layers.
activations (Sequence[str]):
    Sequence of strings containing the activation functions.
train (bool):
    Whether the model should be set to training mode or not.
train_output_type (str):
    The output type of the model during training.

Methods:

• load_state_from_file –

  Loads the state dictionary from a file.

• make_forward_partial –

  Creates a partial function for the forward pass of the network.

• setup –

  Setup function for the JaxMLP class.

load_state_from_file

load_state_from_file(seed: int = 42, input_dim: int = 6, file_path: str | None = None) -> flax.core.frozen_dict.FrozenDict

Loads the state dictionary from a file.

Arguments

seed (int):
    Seed for the random number generator.
input_dim (int):
    Dimension of the input tensor.
file_path (str):
    Path to the file containing the state dictionary.

Returns:

• flax.core.frozen_dict.FrozenDict: –

  The state dictionary.

make_forward_partial

make_forward_partial(seed: int = 42, input_dim: int = 6, state: str | dict | None = None, add_jitted: bool = False) -> tuple[Callable, Callable | None]

Creates a partial function for the forward pass of the network.

Arguments

seed (int):
    Seed for the random number generator.
input_dim (int):
    Dimension of the input tensor.
state (flax.core.frozen_dict.FrozenDict):
    The state dictionary (if not loaded from file).
add_jitted (bool):
    Whether the partial function should be jitted or not.

Returns:

• Callable: –

  The partial function for the forward pass of the network.

setup

setup() -> None

Setup function for the JaxMLP class. Initializes the layers and activation functions.

lanfactory.trainers.LoadTorchMLP

LoadTorchMLP(model_file_path: str, network_config: dict | str, input_dim: int, network_type: str | None = None, inference_mode: bool = False)

General-purpose class to load TorchMLP models.

Does NOT call eval() by default - suitable for fine-tuning or further training. For inference with eval() enabled, use LoadTorchMLPInfer instead.

Arguments

model_file_path (str):
    Path to the model state dict file.
network_config (dict | str):
    Network configuration dictionary or path to a pickled config file.
input_dim (int):
    Input dimension.
network_type (str | None):
    Network type ("lan", "cpn", "opn"). If not provided, will be
    inferred from train_output_type in network_config.
inference_mode (bool):
    If True, sets network to eval mode. Default is False for general use.
    Use LoadTorchMLPInfer for inference with eval() enabled by default.

Methods:

• predict_on_batch –

  Make predictions on a batch of data.

predict_on_batch

predict_on_batch(x: ndarray | None = None) -> np.ndarray

Make predictions on a batch of data.

This method is intended for computing trial-wise log-likelihoods from a matrix input, and is commonly used through the HDDM toolbox.

Arguments

x (numpy.ndarray):
    Input matrix (dtype should be numpy.float32).
    For LANs, columns should follow a specific order:
    model parameters followed by reaction times and choices.

Returns:

• numpy.ndarray: –

  Network output as numpy array.

lanfactory.trainers.LoadTorchMLPInfer

LoadTorchMLPInfer(model_file_path: str, network_config: dict | str, input_dim: int, network_type: str | None = None, inference_mode: bool = True)

Bases: LoadTorchMLP

Model loader with inference mode enabled by default.

Calls eval() on the network, suitable for inference/prediction. For fine-tuning or further training, use LoadTorchMLP instead.

This class was originally useful directly for application in the HDDM toolbox.

Arguments

model_file_path (str):
    Path to the model state dict file.
network_config (dict | str):
    Network configuration dictionary or path to a pickled config file.
input_dim (int):
    Input dimension.
network_type (str | None):
    Network type ("lan", "cpn", "opn"). If not provided, will be
    inferred from train_output_type in network_config.
inference_mode (bool):
    If True, sets network to eval mode. Default is True for inference.

lanfactory.trainers.ModelTrainerJaxMLP

ModelTrainerJaxMLP(train_config: dict, model: JaxMLP, train_dl: Any, valid_dl: Any, allow_abs_path_folder_generation: bool = False, pin_memory: bool = False, seed: int | None = None)

Class for training JaxMLP models.

Arguments

train_config (dict):
    Dictionary containing the training configuration.
model (JaxMLP):
    The JaxMLP model to be trained.
train_dl (torch.utils.data.DataLoader):
    The training data loader.
valid_dl (torch.utils.data.DataLoader):
    The validation data loader.
allow_abs_path_folder_generation (bool):
    Whether the folder for the output files should be created or not.
pin_memory (bool):
    Whether the data loader should pin memory or not.
seed (int):
    Seed for the random number generator.

Returns:

• ModelTrainerJaxMLP: –

  The ModelTrainerJaxMLP object.

Methods:

• create_train_state –

  Create initial train state

• run_epoch –

  Run one epoch of training or validation

• train_and_evaluate –

  Train and evaluate JAXMLP model.

create_train_state

create_train_state(rng: PRNGKey) -> train_state.TrainState

Create initial train state

run_epoch

run_epoch(state: TrainState, train: bool = True, verbose: int = 1, epoch: int = 0, max_epochs: int = 0) -> tuple[train_state.TrainState, float]

Run one epoch of training or validation

Arguments

state (flax.core.frozen_dict.FrozenDict):
    The state dictionary.
train (bool):
    Whether the model should is in training mode or not.
verbose (int):
    The verbosity level.
epoch (int):
    The current epoch.
max_epochs (int):
    The maximum number of epochs.

Returns:

• tuple (flax.core.frozen_dict.FrozenDict, float): –

  The state dictionary and the mean epoch loss.

train_and_evaluate

train_and_evaluate(output_folder: str | Path = 'data/', output_file_id: str = 'fileid', run_id: str = 'runid', mlflow_on: bool = False, save_outputs: bool = True, verbose: int = 1, network_type: str | None = None, export_onnx: bool = True) -> train_state.TrainState

Train and evaluate JAXMLP model.

Arguments

output_folder (str):
    Path to the output folder.
output_file_id (str):
    The file id.
run_id (str):
    The run id.
mlflow_on (bool):
    Whether to use mlflow or not.
save_outputs (bool):
    Whether to save all files or not.
verbose (int):
    The verbosity level.
network_type (str | None):
    The network type ('lan', 'cpn', 'opn', ...), used in output
    filenames. When None it is inferred from train_output_type —
    which cannot distinguish cpn from opn (both use logits), so
    callers that know the type should pass it.
export_onnx (bool):
    Whether to export the trained network to ONNX alongside the
    flax state (single-trial contract; the artifact HSSM consumes).

Returns:

• flax.core.frozen_dict.FrozenDict: –

  The final state dictionary (model state).

lanfactory.trainers.ModelTrainerTorchMLP

ModelTrainerTorchMLP(model: TorchMLP, train_config: dict | str | Path, train_dl: DataLoader, valid_dl: DataLoader, allow_abs_path_folder_generation: bool = False, pin_memory: bool = True, seed: int | None = None)

---

train_config (dict):
    Training configuration.
model (TorchMLP):
    TorchMLP model.
train_dl (DatasetTorch):
    Training dataloader.
valid_dl (DatasetTorch):
    Validation dataloader.
allow_abs_path_folder_generation (bool):
    Whether to allow absolute path folder generation.
pin_memory (bool):
    Whether to pin memory (dataloader). Can affect speed.
seed (int):
    Random seed.

Methods:

• train_and_evaluate –

  Train and evaluate the model.

train_and_evaluate

train_and_evaluate(output_folder: str | Path = 'data/', output_file_id: str = 'fileid', run_id: str = 'runid', mlflow_on: bool = False, save_outputs: bool = True, verbose: int = 1) -> None

Train and evaluate the model.

Arguments

output_folder (str):
    Output folder.
output_file_id (str):
    Output file ID.
run_id (str):
    Run ID.
mlflow_on (bool):
    Whether to use mlflow.
save_outputs (bool):
    Whether to save all outputs.
verbose (int):
    Verbosity level.

lanfactory.trainers.TorchMLP

TorchMLP(network_config: dict, input_shape: int = 10, network_type: str | None = None)

Bases: Module

TorchMLP class.

Arguments

network_config (dict):
    Network configuration.
input_shape (int):
    Input shape.
network_type (str):
    Network type.

Methods:

• forward –

  Forward pass through network.

forward

forward(x: Tensor) -> torch.Tensor

Forward pass through network.

Arguments

x (torch.Tensor):
    Input tensor.

Returns:

• torch.Tensor: –

  Output tensor.

lanfactory.trainers.JaxMLPFactory

JaxMLPFactory(network_config: dict | str | None = None, train: bool = True) -> JaxMLP

Factory function to create a JaxMLP object.

Arguments

network_config (dict | str | None):
    Dictionary containing the network configuration or path to pickled config.
train (bool):
    Whether the model should be trained or not.

Returns:

• JaxMLP class initialized with the correct network configuration. –

lanfactory.trainers.TorchMLPFactory

TorchMLPFactory(network_config: dict | str, input_dim: int, network_type: str | None = None) -> TorchMLP

Factory function to create a TorchMLP object.

This provides a consistent API with JaxMLPFactory and handles loading network configs from pickle files.

Arguments

network_config: Dictionary containing the network configuration,
    or path to a pickled config file.
input_dim: Input dimension (typically from dataloader.dataset.input_dim).
network_type: Network type ("lan", "cpn", "opn"). If not provided,
    will be inferred from train_output_type in network_config.

Returns:

• TorchMLP instance ready for training. –

Example

>>> train_dl, valid_dl, input_dim = make_train_valid_dataloaders(...)
>>> net = TorchMLPFactory(
...     network_config=network_config,
...     input_dim=input_dim,
... )

lanfactory.trainers.make_dataloader

make_dataloader(file_ids: list[str] | list[Path], batch_size: int, network_type: str = 'lan', label_lower_bound: float | None = None, shuffle: bool = True, num_workers: int = 1, pin_memory: bool = True) -> DataLoader

Create a DataLoader for LAN/CPN/OPN training.

This is a convenience function that creates a DatasetTorch and wraps it in a PyTorch DataLoader with sensible defaults.

Arguments

file_ids: List of paths to training data pickle files.
batch_size: Batch size for training.
network_type: Type of network ("lan", "cpn", or "opn").
    Determines the feature/label keys in the data files.
label_lower_bound: Lower bound for labels. If None and network_type
    is "lan", defaults to log(1e-10).
shuffle: Whether to shuffle data (default: True).
num_workers: Number of worker processes for data loading (default: 1).
pin_memory: Whether to pin memory for faster GPU transfer (default: True).

Returns:

• torch.utils.data.DataLoader configured for training. –

Example

>>> file_list = list(Path("data/lan_mlp/ddm").glob("*.pickle"))
>>> train_dl = make_dataloader(
...     file_ids=file_list,
...     batch_size=4096,
...     network_type="lan",
... )

lanfactory.trainers.make_train_valid_dataloaders

make_train_valid_dataloaders(file_ids: list[str] | list[Path], batch_size: int, network_type: str = 'lan', train_val_split: float = 0.9, shuffle_files: bool = True, label_lower_bound: float | None = None, num_workers: int = 1, pin_memory: bool = True) -> tuple[DataLoader, DataLoader, int]

Create train and validation DataLoaders with automatic file splitting.

This is a convenience function that splits the file list into train/validation sets and creates DataLoaders for each.

Arguments

file_ids: List of paths to training data pickle files.
batch_size: Batch size for training.
network_type: Type of network ("lan", "cpn", or "opn").
train_val_split: Fraction of files to use for training (default: 0.9).
shuffle_files: Whether to shuffle files before splitting (default: True).
label_lower_bound: Lower bound for labels. If None and network_type
    is "lan", defaults to log(1e-10).
num_workers: Number of worker processes for data loading (default: 1).
pin_memory: Whether to pin memory for faster GPU transfer (default: True).

Returns:

• tuple of (train_dataloader, valid_dataloader, input_dim) –

Example

>>> file_list = list(Path("data/lan_mlp/ddm").glob("*.pickle"))
>>> train_dl, valid_dl, input_dim = make_train_valid_dataloaders(
...     file_ids=file_list,
...     batch_size=4096,
...     network_type="lan",
...     train_val_split=0.9,
... )