braintools.trainer module

Trainer module for PyTorch Lightning-like training in JAX.

This module provides a comprehensive training framework including: - LightningModule: Base class for defining training models - Trainer: Orchestration class for training loops - Callbacks: Hook system for customizing training behavior - Loggers: Pluggable logging backends (TensorBoard, WandB, CSV, etc.) - DataLoader: JAX-compatible data loading with distributed support - Distributed: Strategies for multi-device and multi-host training

Example

>>> import braintools
>>> import brainstate
>>>
>>> class MyModel(braintools.trainer.LightningModule):
...     def __init__(self):
...         super().__init__()
...         self.linear = brainstate.nn.Linear(784, 10)
...
...     def __call__(self, x):
...         return self.linear(x)
...
...     def training_step(self, batch, batch_idx):
...         x, y = batch['x'], batch['y']
...         logits = self(x)
...         loss = jnp.mean((logits - y) ** 2)
...         self.log('train_loss', loss)
...         return {'loss': loss}
...
...     def configure_optimizers(self):
...         return braintools.optim.Adam(lr=1e-3)
...
>>> model = MyModel()
>>> trainer = braintools.trainer.Trainer(max_epochs=10)
>>> train_loader = braintools.trainer.DataLoader(data, batch_size=32)
>>> trainer.fit(model, train_loader)

A PyTorch-Lightning-style training framework for JAX / brainstate models. It bundles the training loop, validation/testing/prediction, a callback and logging system, checkpointing, data loading, and multi-device strategies into a single high-level API.

Overview

The braintools.trainer module provides:

• LightningModule – base class for defining models, training/validation steps, and optimizers.

• Trainer – orchestrates the full fit/validate/test/predict loops.

• Callbacks – hooks for checkpointing, early stopping, LR monitoring, etc.

• Loggers – pluggable logging backends (CSV, TensorBoard, WandB, …).

• DataLoader – JAX-compatible batching, shuffling, and sampling.

• Distributed strategies – single-device, data-parallel, sharded, and FSDP.

• Checkpointing – save/restore model and optimizer state.

Note

Metrics logged via LightningModule.log() are referenced by the trainer and callbacks under the name you pass. The trainer prefixes aggregated validation/test metrics with val_/test_ only when the name does not already start with that prefix, so self.log('val_loss', ...) is surfaced as val_loss (not val_val_loss).

Core Classes

The two classes you interact with most: subclass LightningModule to define your model and call Trainer.fit() to train it.

LightningModule	Base class for all training modules.
Trainer	Orchestrates the training process.
TrainerState	Container for trainer state during training.
TrainOutput	Output container for training step.
EvalOutput	Output container for evaluation steps (validation/test).

Callbacks

Callbacks hook into the training loop to add behavior without modifying the model. Callback is the base class; CallbackList dispatches to a list of callbacks.

Callback	Base class for callbacks.
CallbackList	Container for multiple callbacks.
ModelCheckpoint	Save model checkpoints based on monitored metric.
EarlyStopping	Stop training when a monitored metric has stopped improving.
LearningRateMonitor	Monitor and log learning rate during training.
GradientClipCallback	Placeholder callback for gradient clipping configuration.
Timer	Track training time.
RichProgressBar	Rich progress bar for training visualization.
TQDMProgressBar	TQDM progress bar for training visualization.
LambdaCallback	Create a callback from lambda functions.
PrintCallback	Simple callback that prints training progress.

Loggers

Pluggable logging backends. Pass an instance (or list) to the logger argument of Trainer.

Logger	Abstract base class for all loggers.
CSVLogger	CSV file logging backend.
TensorBoardLogger	TensorBoard logging backend.
WandBLogger	Weights & Biases logging backend.
NeptuneLogger	Neptune.ai logging backend.
MLFlowLogger	MLflow logging backend.
CompositeLogger	Combine multiple loggers.

Data Loading

JAX-compatible datasets, samplers, and loaders.

DataLoader	JAX-compatible data loader.
DistributedDataLoader	DataLoader for distributed training.
Dataset	Abstract base class for datasets.
ArrayDataset	Dataset wrapping arrays.
DictDataset	Dataset wrapping a dictionary of arrays.
IterableDataset	Dataset wrapping an iterable.
Sampler	Base class for samplers.
RandomSampler	Samples elements randomly.
SequentialSampler	Samples elements sequentially.
BatchSampler	Wraps a sampler to yield batches of indices.
DistributedSampler	Sampler for distributed training.

create_distributed_batches

Create batches suitable for distributed training with pmap.

Distributed Strategies

Strategies that control how the model and data are distributed across devices.

Strategy	Abstract base class for distributed training strategies.
SingleDeviceStrategy	Single device strategy (no distribution).
DataParallelStrategy	Data parallel strategy using pmap.
ShardedDataParallelStrategy	Sharded data parallel strategy using jax.sharding.
FullyShardedDataParallelStrategy	Fully Sharded Data Parallel (FSDP) strategy.
AutoStrategy	Automatically select the best strategy based on available devices.

get_strategy	Get a distributed training strategy by name.
all_reduce	All-reduce tensor across devices.
broadcast	Broadcast tensor from source device to all devices.

Checkpointing

Utilities to save and restore training state.

CheckpointManager

Manager for saving and loading checkpoints.

save_checkpoint	Save a checkpoint to disk.
load_checkpoint	Load a checkpoint from disk.
find_checkpoint	Find a checkpoint file in a directory.
list_checkpoints	List all checkpoints in a directory.

Progress Bars

ProgressBar	Abstract base class for progress bars.
SimpleProgressBar	Simple text-based progress bar with no dependencies.
TQDMProgressBarWrapper	Progress bar wrapper using tqdm.
RichProgressBarWrapper	Progress bar wrapper using rich.

get_progress_bar

Get a progress bar instance.