In [ ]:

import marimo as mo

import marimo as mo

Train with the JAX backend¶

This is the JAX/Flax companion to Train your first LAN (PyTorch). Complete that learning tutorial first: it owns the explanation of LAN training data, configuration, and the end-to-end workflow. Here we repeat a deliberately tiny data fixture so this how-to remains executable, then focus on the JAX-specific factory, trainer, saved state, and inference call.

The fixture uses ssm-simulators to generate DDM data. LANfactory can train on compatible data from other generators as well.

In [ ]:

from copy import deepcopy
from pathlib import Path

import lanfactory
import numpy as np
import ssms

from copy import deepcopy from pathlib import Path import lanfactory import numpy as np import ssms

Generate a small training fixture¶

These settings make two small files so the example runs quickly. For real training, choose the simulator budget and parameter coverage described in the PyTorch learning tutorial and the ssm-simulators data-generation documentation.

In [ ]:

MODEL = "ddm"
OUT_FOLDER = Path("jax_nb_data") / "training_data"
MODEL_FOLDER = Path("jax_nb_data") / "jax_models" / "lan"
N_DATA_FILES = 2
BATCH_SIZE = 1000
OUT_FOLDER.mkdir(parents=True, exist_ok=True)
MODEL_FOLDER.mkdir(parents=True, exist_ok=True)

generator_config = ssms.config.get_default_generator_config("lan")
generator_config["model"] = MODEL
generator_config["pipeline"]["n_parameter_sets"] = 100
generator_config["pipeline"]["n_cpus"] = 1
generator_config["simulator"]["n_samples"] = 200
generator_config["training"]["n_samples_per_param"] = 200
generator_config["output"]["folder"] = str(OUT_FOLDER)

model_config = deepcopy(ssms.config.model_config[MODEL])

MODEL = "ddm" OUT_FOLDER = Path("jax_nb_data") / "training_data" MODEL_FOLDER = Path("jax_nb_data") / "jax_models" / "lan" N_DATA_FILES = 2 BATCH_SIZE = 1000 OUT_FOLDER.mkdir(parents=True, exist_ok=True) MODEL_FOLDER.mkdir(parents=True, exist_ok=True) generator_config = ssms.config.get_default_generator_config("lan") generator_config["model"] = MODEL generator_config["pipeline"]["n_parameter_sets"] = 100 generator_config["pipeline"]["n_cpus"] = 1 generator_config["simulator"]["n_samples"] = 200 generator_config["training"]["n_samples_per_param"] = 200 generator_config["output"]["folder"] = str(OUT_FOLDER) model_config = deepcopy(ssms.config.model_config[MODEL])

In [ ]:

for _i in range(N_DATA_FILES):
    print(f"Generating data file {_i + 1} / {N_DATA_FILES}")
    _generator = ssms.dataset_generators.lan_mlp.TrainingDataGenerator(
        config=generator_config,
        model_config=model_config,
    )
    _generator.generate_data_training(save=True)

for _i in range(N_DATA_FILES): print(f"Generating data file {_i + 1} / {N_DATA_FILES}") _generator = ssms.dataset_generators.lan_mlp.TrainingDataGenerator( config=generator_config, model_config=model_config, ) _generator.generate_data_training(save=True)

Generating data file 1 / 2
Generating data file 2 / 2

In [ ]:

network_config = deepcopy(lanfactory.config.network_configs.network_config_mlp)
network_config["layer_sizes"] = [100, 100, 100, 1]
network_config["activations"] = ["tanh", "tanh", "tanh", "linear"]
print("Network config:")
print(network_config)

train_config = deepcopy(lanfactory.config.network_configs.train_config_mlp)
train_config["learning_rate"] = 2e-6
train_config["cpu_batch_size"] = 4096
train_config["gpu_batch_size"] = 4096
train_config["n_epochs"] = 2
print("Train config:")
print(train_config)

network_config = deepcopy(lanfactory.config.network_configs.network_config_mlp) network_config["layer_sizes"] = [100, 100, 100, 1] network_config["activations"] = ["tanh", "tanh", "tanh", "linear"] print("Network config:") print(network_config) train_config = deepcopy(lanfactory.config.network_configs.train_config_mlp) train_config["learning_rate"] = 2e-6 train_config["cpu_batch_size"] = 4096 train_config["gpu_batch_size"] = 4096 train_config["n_epochs"] = 2 print("Train config:") print(train_config)

Network config:
{'layer_sizes': [100, 100, 100, 1], 'activations': ['tanh', 'tanh', 'tanh', 'linear'], 'train_output_type': 'logprob'}
Train config:
{'cpu_batch_size': 4096, 'gpu_batch_size': 4096, 'n_epochs': 2, 'optimizer': 'adam', 'learning_rate': 2e-06, 'lr_scheduler': 'reduce_on_plateau', 'lr_scheduler_params': {}, 'weight_decay': 0.0, 'loss': 'huber', 'save_history': True}

Reuse the shared data loader¶

make_train_valid_dataloaders is backend-neutral at this boundary. It:

• splits your data files into training and validation sets;
• creates the appropriate DatasetTorch objects; and
• wraps them in PyTorch DataLoader objects with sensible defaults.

The loaders yield NumPy-compatible batches that the JAX trainer consumes.

In [ ]:

file_list_ = sorted(OUT_FOLDER.glob("*.pickle"))

# num_workers=0 keeps data loading in-process: ssm-simulators sets the
# multiprocessing start method to "spawn" at import, which is unsafe for
# DataLoader worker processes inside a notebook.
jax_training_dataloader, jax_validation_dataloader, input_dim = lanfactory.trainers.make_train_valid_dataloaders(
    file_ids=file_list_,
    batch_size=BATCH_SIZE,
    network_type="lan",
    train_val_split=0.5,
    num_workers=0,
    pin_memory=False,
)

print(f"Training batches: {len(jax_training_dataloader)}")
print(f"Validation batches: {len(jax_validation_dataloader)}")
print(f"Input dimension: {input_dim}")

file_list_ = sorted(OUT_FOLDER.glob("*.pickle")) # num_workers=0 keeps data loading in-process: ssm-simulators sets the # multiprocessing start method to "spawn" at import, which is unsafe for # DataLoader worker processes inside a notebook. jax_training_dataloader, jax_validation_dataloader, input_dim = lanfactory.trainers.make_train_valid_dataloaders( file_ids=file_list_, batch_size=BATCH_SIZE, network_type="lan", train_val_split=0.5, num_workers=0, pin_memory=False, ) print(f"Training batches: {len(jax_training_dataloader)}") print(f"Validation batches: {len(jax_validation_dataloader)}") print(f"Input dimension: {input_dim}")

Training batches: 20
Validation batches: 20
Input dimension: 6

Create the JAX/Flax network¶

JaxMLPFactory materializes the configured Flax MLP. Passing train=True prepares it for the training-state lifecycle used by ModelTrainerJaxMLP.

In [ ]:

jax_net = lanfactory.trainers.JaxMLPFactory(
    network_config=network_config,
    train=True,
)

jax_net = lanfactory.trainers.JaxMLPFactory( network_config=network_config, train=True, )

Train and save the JAX state¶

The JAX trainer takes the same configuration and loaders as the PyTorch path, but writes a Flax training-state artifact rather than a Torch model.

In [ ]:

jax_trainer = lanfactory.trainers.ModelTrainerJaxMLP(
    train_config=train_config,
    model=jax_net,
    train_dl=jax_training_dataloader,
    valid_dl=jax_validation_dataloader,
    pin_memory=False,
)

jax_trainer = lanfactory.trainers.ModelTrainerJaxMLP( train_config=train_config, model=jax_net, train_dl=jax_training_dataloader, valid_dl=jax_validation_dataloader, pin_memory=False, )

In [ ]:

train_state = jax_trainer.train_and_evaluate(
    output_folder=MODEL_FOLDER,
    output_file_id=MODEL,
    run_id="jax",
    mlflow_on=False,
    verbose=1,
    save_outputs=True,
)

train_state = jax_trainer.train_and_evaluate( output_folder=MODEL_FOLDER, output_file_id=MODEL, run_id="jax", mlflow_on=False, verbose=1, save_outputs=True, )

Epoch: 0 of 2
Training - Step: 0 of 20 - Loss: 4.5220714
Epoch 0/2 time: 0.21618890762329102s
Validation - Step: 0 of 20 - Loss: 3.173557
Epoch 0/2 time: 0.06603813171386719s
Epoch: 0 / 2, test_loss: 3.1969714164733887
Epoch: 1 of 2
Training - Step: 0 of 20 - Loss: 2.8688884
Epoch 1/2 time: 0.0069768428802490234s
Validation - Step: 0 of 20 - Loss: 1.7008276
Epoch 1/2 time: 0.0024831295013427734s
Epoch: 1 / 2, test_loss: 1.8143354654312134
Saving training history to: jax_nb_data/jax_models/lan/jax_lan_ddm__jax_training_history.csv
Saving model parameters to: jax_nb_data/jax_models/lan/jax_lan_ddm__train_state.jax
Saving training config to: jax_nb_data/jax_models/lan/jax_lan_ddm__train_config.pickle
Saving training data details to: jax_nb_data/jax_models/lan/jax_lan_ddm__data_details.pickle
Saving ONNX export to: jax_nb_data/jax_models/lan/jax_lan_ddm__model.onnx

Reload the state for inference¶

Recreate the factory with train=False, load the saved state, and request a JIT-compiled forward function. The input width is the model parameter count plus reaction time and response.

In [ ]:

jax_infer = lanfactory.trainers.JaxMLPFactory(
    network_config=network_config,
    train=False,
)

jax_infer = lanfactory.trainers.JaxMLPFactory( network_config=network_config, train=False, )

In [ ]:

# Establish the reactive dependency before reloading the saved state.
_ = train_state
_forward_pass, forward_pass_jitted = jax_infer.make_forward_partial(
    seed=42,
    input_dim=model_config["n_params"] + 2,
    state=str(MODEL_FOLDER / f"jax_lan_{MODEL}__train_state.jax"),
    add_jitted=True,
)

# Establish the reactive dependency before reloading the saved state. _ = train_state _forward_pass, forward_pass_jitted = jax_infer.make_forward_partial( seed=42, input_dim=model_config["n_params"] + 2, state=str(MODEL_FOLDER / f"jax_lan_{MODEL}__train_state.jax"), add_jitted=True, )

In [ ]:

import jax.numpy as jnp

theta = deepcopy(ssms.config.model_config[MODEL]["default_params"])
sim_out = ssms.basic_simulators.simulator.simulator(
    model=MODEL,
    theta=theta,
    n_samples=50_000,
)
input_mat = jnp.zeros((2000, len(theta) + 2))
for _i in range(len(theta)):
    input_mat = input_mat.at[:, _i].set(jnp.ones(2000) * theta[_i])
input_mat = input_mat.at[:, len(theta)].set(
    jnp.array(
        np.concatenate(
            [
                np.linspace(5, 0, 1000, dtype=np.float32),
                np.linspace(0, 5, 1000, dtype=np.float32),
            ]
        )
    )
)
input_mat = input_mat.at[:, len(theta) + 1].set(
    jnp.array(
        np.concatenate(
            [np.repeat(-1.0, 1000), np.repeat(1.0, 1000)]
        ).astype(np.float32)
    )
)
net_out = forward_pass_jitted(input_mat)

import jax.numpy as jnp theta = deepcopy(ssms.config.model_config[MODEL]["default_params"]) sim_out = ssms.basic_simulators.simulator.simulator( model=MODEL, theta=theta, n_samples=50_000, ) input_mat = jnp.zeros((2000, len(theta) + 2)) for _i in range(len(theta)): input_mat = input_mat.at[:, _i].set(jnp.ones(2000) * theta[_i]) input_mat = input_mat.at[:, len(theta)].set( jnp.array( np.concatenate( [ np.linspace(5, 0, 1000, dtype=np.float32), np.linspace(0, 5, 1000, dtype=np.float32), ] ) ) ) input_mat = input_mat.at[:, len(theta) + 1].set( jnp.array( np.concatenate( [np.repeat(-1.0, 1000), np.repeat(1.0, 1000)] ).astype(np.float32) ) ) net_out = forward_pass_jitted(input_mat)

In [ ]:

from matplotlib import pyplot as plt

plt.plot(
    np.linspace(-5, 5, 2000, dtype=np.float32),
    np.exp(net_out),
    label="JAX LAN likelihood",
)
plt.hist(
    sim_out["rts"] * sim_out["choices"],
    bins=100,
    histtype="step",
    fill=None,
    density=True,
    label="simulated choices × RT",
)
plt.legend()
plt.show()

from matplotlib import pyplot as plt plt.plot( np.linspace(-5, 5, 2000, dtype=np.float32), np.exp(net_out), label="JAX LAN likelihood", ) plt.hist( sim_out["rts"] * sim_out["choices"], bins=100, histtype="step", fill=None, density=True, label="simulated choices × RT", ) plt.legend() plt.show()