QML tools

ML Tools

This module implements a Trainer class for torch Modules and QuantumModel. It also implements the QNN class and callbacks that can be used with the trainer module.

`Trainer(model, optimizer, config, loss_fn='mse', train_dataloader=None, val_dataloader=None, test_dataloader=None, optimize_step=optimize_step, max_batches=None)`

Bases: BaseTrainer

Trainer class to manage and execute training, validation, and testing loops for a model (eg.

QNN).

This class handles the overall training process, including: - Managing epochs and steps - Handling data loading and batching - Computing and updating gradients - Logging and monitoring training metrics

ATTRIBUTE	DESCRIPTION
`current_epoch`	The current epoch number. TYPE: `int`
`global_step`	The global step across all epochs. TYPE: `int`

Inherited Attributes

use_grad (bool): Indicates if gradients are used for optimization. Default is True.

model (nn.Module): The neural network model. optimizer (optim.Optimizer | NGOptimizer | None): The optimizer for training. config (TrainConfig): The configuration settings for training. train_dataloader (DataLoader | DictDataLoader | None): DataLoader for training data. val_dataloader (DataLoader | DictDataLoader | None): DataLoader for validation data. test_dataloader (DataLoader | DictDataLoader | None): DataLoader for testing data.

optimize_step (Callable): Function for performing an optimization step. loss_fn (Callable): loss function to use.

num_training_batches (int): Number of training batches. num_validation_batches (int): Number of validation batches. num_test_batches (int): Number of test batches.

state (str): Current state in the training process

Default training routine

for epoch in max_iter + 1:
    # Training
    for batch in train_batches:
        train model
    # Validation
    if val_every % epoch == 0:
        for batch in val_batches:
            train model

Notes

In case of InfiniteTensorDataset, number of batches = 1.
In case of TensorDataset, number of batches are default.
Training is run for max_iter + 1 epochs. Epoch 0 logs untrained model.
Please look at the CallbackManager initialize_callbacks method to review the default logging behavior.

Examples:

import torch
from torch.optim import SGD
from qadence import (
    feature_map,
    hamiltonian_factory,
    hea,
    QNN,
    QuantumCircuit,
    TrainConfig,
    Z,
)
from qadence.ml_tools.trainer import Trainer
from qadence.ml_tools.optimize_step import optimize_step
from qadence.ml_tools import TrainConfig
from qadence.ml_tools.data import to_dataloader

# Initialize the model
n_qubits = 2
fm = feature_map(n_qubits)
ansatz = hea(n_qubits=n_qubits, depth=2)
observable = hamiltonian_factory(n_qubits, detuning=Z)
circuit = QuantumCircuit(n_qubits, fm, ansatz)
model = QNN(circuit, observable, backend="pyqtorch", diff_mode="ad")

# Set up the optimizer
optimizer = SGD(model.parameters(), lr=0.001)

# Use TrainConfig for configuring the training process
config = TrainConfig(
    max_iter=100,
    print_every=10,
    write_every=10,
    checkpoint_every=10,
    val_every=10
)

# Create the Trainer instance with TrainConfig
trainer = Trainer(
    model=model,
    optimizer=optimizer,
    config=config,
    loss_fn="mse",
    optimize_step=optimize_step
)

batch_size = 25
x = torch.linspace(0, 1, 32).reshape(-1, 1)
y = torch.sin(x)
train_loader = to_dataloader(x, y, batch_size=batch_size, infinite=True)
val_loader = to_dataloader(x, y, batch_size=batch_size, infinite=False)

# Train the model
model, optimizer = trainer.fit(train_loader, val_loader)

This also supports both gradient based and gradient free optimization. The default support is for gradient based optimization.

Notes:

set_use_grad() (class level):This method is used to set the global use_grad flag, controlling whether the trainer uses gradient-based optimization.
```
# gradient based
Trainer.set_use_grad(True)

# gradient free
Trainer.set_use_grad(False)
```
Context Managers (instance level): enable_grad_opt() and disable_grad_opt() are context managers that temporarily switch the optimization mode for specific code blocks. This is useful when you want to mix gradient-based and gradient-free optimization in the same training process.
```
# gradient based
with trainer.enable_grad_opt(optimizer):
    trainer.fit()

# gradient free
with trainer.disable_grad_opt(ng_optimizer):
    trainer.fit()
```

Examples

Gradient based optimization example Usage:

from torch import optim
optimizer = optim.SGD(model.parameters(), lr=0.01)

Trainer.set_use_grad(True)
trainer = Trainer(
    model=model,
    optimizer=optimizer,
    config=config,
    loss_fn="mse"
)
trainer.fit(train_loader, val_loader)

or

trainer = Trainer(
    model=model,
    config=config,
    loss_fn="mse"
)
with trainer.enable_grad_opt(optimizer):
    trainer.fit(train_loader, val_loader)

Gradient free optimization example Usage:

import nevergrad as ng
from qadence.ml_tools.parameters import num_parameters
ng_optimizer = ng.optimizers.NGOpt(
                budget=config.max_iter, parametrization= num_parameters(model)
                )

Trainer.set_use_grad(False)
trainer = Trainer(
    model=model,
    optimizer=ng_optimizer,
    config=config,
    loss_fn="mse"
)
trainer.fit(train_loader, val_loader)

or

import nevergrad as ng
from qadence.ml_tools.parameters import num_parameters
ng_optimizer = ng.optimizers.NGOpt(
        budget=config.max_iter, parametrization= num_parameters(model)
        )

trainer = Trainer(
    model=model,
    config=config,
    loss_fn="mse"
)
with trainer.disable_grad_opt(ng_optimizer):
    trainer.fit(train_loader, val_loader)

Initializes the Trainer class.

PARAMETER	DESCRIPTION
`model`	The PyTorch model to train. TYPE: `Module`
`optimizer`	The optimizer for training. TYPE: `Optimizer \| Optimizer \| None`
`config`	Training configuration object. TYPE: `TrainConfig`
`loss_fn`	Loss function used for training. If not specified, default mse loss will be used. TYPE: `str \| Callable` DEFAULT: `'mse'`
`train_dataloader`	DataLoader for training data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`val_dataloader`	DataLoader for validation data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`test_dataloader`	DataLoader for test data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`optimize_step`	Function to execute an optimization step. TYPE: `Callable` DEFAULT: `optimize_step`
`max_batches`	Maximum number of batches to process per epoch. This is only valid in case of finite TensorDataset dataloaders. if max_batches is not None, the maximum number of batches used will be min(max_batches, len(dataloader.dataset)) In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: `int \| None` DEFAULT: `None`

Source code in qadence/ml_tools/trainer.py

def __init__(
    self,
    model: nn.Module,
    optimizer: optim.Optimizer | NGOptimizer | None,
    config: TrainConfig,
    loss_fn: str | Callable = "mse",
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
    test_dataloader: DataLoader | DictDataLoader | None = None,
    optimize_step: Callable = optimize_step,
    max_batches: int | None = None,
):
    """
    Initializes the Trainer class.

    Args:
        model (nn.Module): The PyTorch model to train.
        optimizer (optim.Optimizer | NGOptimizer | None): The optimizer for training.
        config (TrainConfig): Training configuration object.
        loss_fn (str | Callable ): Loss function used for training.
            If not specified, default mse loss will be used.
        train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
        val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.
        test_dataloader (DataLoader | DictDataLoader |  None): DataLoader for test data.
        optimize_step (Callable): Function to execute an optimization step.
        max_batches (int | None): Maximum number of batches to process per epoch.
            This is only valid in case of finite TensorDataset dataloaders.
            if max_batches is not None, the maximum number of batches used will
            be min(max_batches, len(dataloader.dataset))
            In case of InfiniteTensorDataset only 1 batch per epoch is used.
    """
    super().__init__(
        model=model,
        optimizer=optimizer,
        config=config,
        loss_fn=loss_fn,
        optimize_step=optimize_step,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        test_dataloader=test_dataloader,
        max_batches=max_batches,
    )
    self.current_epoch: int = 0
    self.global_step: int = 0
    self._stop_training: torch.Tensor = torch.tensor(0, dtype=torch.int)
    self.progress: Progress | None = None

    # Integration with Accelerator:
    self.accelerator = Accelerator(
        backend=config.backend,
        nprocs=config.nprocs,
        compute_setup=config.compute_setup,
        dtype=config.dtype,
        log_setup=config.log_setup,
    )
    # Decorate the unbound Trainer.fit method with accelerator.distribute.
    # We use __get__ to bind the decorated method to the current instance,
    # ensuring that 'self' is passed only once when self.fit is called.
    self.fit = self.accelerator.distribute(Trainer.fit).__get__(self, Trainer)  # type: ignore[method-assign]

`_aggregate_result(result)`

Aggregates the loss and metrics using the Accelerator's all_reduce_dict method if aggregation is enabled.

PARAMETER	DESCRIPTION
`result`	(tuple[torch.Tensor, dict[str, Any]]) The result consisting of loss and metrics.For more details, look at the signature of build_optimize_result. TYPE: `tuple[Tensor, dict[str, Any]]`

RETURNS	DESCRIPTION
`tuple[Tensor, dict[str, Any]]`	tuple[torch.Tensor, dict[str, Any]]: The aggregated loss and metrics.

Source code in qadence/ml_tools/trainer.py

def _aggregate_result(
    self, result: tuple[torch.Tensor, dict[str, Any]]
) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Aggregates the loss and metrics using the Accelerator's all_reduce_dict method if aggregation is enabled.

    Args:
        result:     (tuple[torch.Tensor, dict[str, Any]])
                        The result consisting of loss and metrics.For more details,
                        look at the signature of build_optimize_result.

    Returns:
        tuple[torch.Tensor, dict[str, Any]]: The aggregated loss and metrics.
    """
    loss, metrics = result
    if self.config.all_reduce_metrics:
        reduced = self.accelerator.all_reduce_dict({"loss": loss, **metrics})
        loss = reduced.pop("loss")
        metrics = reduced
        return loss, metrics
    else:
        return loss, metrics

`_batch_iter(dataloader, num_batches)`

Yields batches from the provided dataloader.

The batch of data is also moved to the correct device and dtype using accelerator.prepare.

PARAMETER	DESCRIPTION
`dataloader`	The dataloader to iterate over. TYPE: `[DataLoader]`
`num_batches`	The maximum number of batches to yield. TYPE: `int`

YIELDS	DESCRIPTION
`Iterable[tuple[Tensor, ...] \| None]`	Iterable[tuple[torch.Tensor, ...] \| None]: A batch from the dataloader moved to the specified device and dtype.

Source code in qadence/ml_tools/trainer.py

def _batch_iter(
    self,
    dataloader: DataLoader | DictDataLoader,
    num_batches: int,
) -> Iterable[tuple[torch.Tensor, ...] | None]:
    """
    Yields batches from the provided dataloader.

    The batch of data is also moved
    to the correct device and dtype using accelerator.prepare.

    Args:
        dataloader ([DataLoader]): The dataloader to iterate over.
        num_batches (int): The maximum number of batches to yield.

    Yields:
        Iterable[tuple[torch.Tensor, ...] | None]: A batch from the dataloader moved to the
            specified device and dtype.
    """
    if dataloader is None:
        for _ in range(num_batches):
            yield None
    else:
        for batch in islice(dataloader, num_batches):
            yield self.accelerator.prepare_batch(batch)

`_fit_end()`

Finalizes the training and closes the writer.

Source code in qadence/ml_tools/trainer.py

def _fit_end(self) -> None:
    """Finalizes the training and closes the writer."""
    self.callback_manager.end_training(trainer=self)

`_fit_setup()`

Sets up the training environment, initializes configurations,.

and moves the model to the specified device and data type. The callback_manager.start_training takes care of loading checkpoint, and setting up the writer.

Source code in qadence/ml_tools/trainer.py

def _fit_setup(self) -> None:
    """
    Sets up the training environment, initializes configurations,.

    and moves the model to the specified device and data type.
    The callback_manager.start_training takes care of loading checkpoint,
    and setting up the writer.
    """
    self._stop_training = torch.tensor(
        0, dtype=torch.int, device=self.accelerator.execution.device
    )
    # initalize config in the first process, and broadcast it to all processes
    if self.accelerator.rank == 0:
        self.config_manager.initialize_config()
    self.config_manager = self.accelerator.broadcast(self.config_manager, src=0)
    self.callback_manager.start_training(trainer=self)

    # Integration with Accelerator: prepare the model, optimizer, and dataloaders.
    (self.model, self.optimizer, self.train_dataloader, self.val_dataloader) = (
        self.accelerator.prepare(
            self.model, self.optimizer, self.train_dataloader, self.val_dataloader
        )
    )

    # Progress bar for training visualization
    if self.accelerator.world_size == 1:
        self.progress = Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
            TimeRemainingColumn(elapsed_when_finished=True),
        )

    # Run validation at the start if specified in the configuration
    self.perform_val = self.config.val_every > 0
    if self.perform_val:
        self.run_validation(self.val_dataloader)

`_modify_batch_end_loss_metrics(loss_metrics)`

Modifies the loss and metrics at the end of batch for proper logging.

All metrics are prefixed with the proper state of the training process - "train_" or "val_" or "test_" A "{state}_loss" is added to metrics.

PARAMETER	DESCRIPTION
`loss_metrics`	Original loss and metrics. TYPE: `tuple[Tensor, dict[str, Any]]`

RETURNS	DESCRIPTION
`tuple[Tensor, dict[str, Any]]`	tuple[None \| torch.Tensor, dict[str, Any]]: Modified loss and metrics.

Source code in qadence/ml_tools/trainer.py

def _modify_batch_end_loss_metrics(
    self, loss_metrics: tuple[torch.Tensor, dict[str, Any]]
) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Modifies the loss and metrics at the end of batch for proper logging.

    All metrics are prefixed with the proper state of the training process
     - "train_" or "val_" or "test_"
    A "{state}_loss" is added to metrics.

    Args:
        loss_metrics (tuple[torch.Tensor, dict[str, Any]]): Original loss and metrics.

    Returns:
        tuple[None | torch.Tensor, dict[str, Any]]: Modified loss and metrics.
    """
    for phase in ["train", "val", "test"]:
        if phase in self.training_stage:
            loss, metrics = loss_metrics
            updated_metrics = {f"{phase}_{key}": value for key, value in metrics.items()}
            updated_metrics[f"{phase}_loss"] = loss
            return loss, updated_metrics
    return loss_metrics

`_train()`

Runs the main training loop over multiple epochs.

This method sets up the training process by performing any necessary pre-training actions (via on_train_start), configuring progress tracking (if available), and then iteratively calling _train_epoch to run through the epochs.

RETURNS	DESCRIPTION
`list[list[tuple[Tensor, dict[str, Any]]]]`	list[list[tuple[torch.Tensor, dict[str, Any]]]]: Training loss
`list[list[tuple[Tensor, dict[str, Any]]]]`	metrics for all epochs. list -> list -> tuples Epochs -> Training Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("train")
def _train(self) -> list[list[tuple[torch.Tensor, dict[str, Any]]]]:
    """
    Runs the main training loop over multiple epochs.

    This method sets up the training process by performing any necessary pre-training
    actions (via `on_train_start`), configuring progress tracking (if available), and then
    iteratively calling `_train_epoch` to run through the epochs.

    Returns:
        list[list[tuple[torch.Tensor, dict[str, Any]]]]: Training loss
        metrics for all epochs.
            list    -> list                  -> tuples
            Epochs  -> Training Batches      -> (loss, metrics)
    """
    self.on_train_start()
    epoch_start, epoch_end = (
        self.global_step,
        self.global_step + self.config_manager.config.max_iter + 1,
    )

    if self.accelerator.world_size == 1 and self.progress:
        # Progress setup is only available for non-spawned training.
        with self.progress:
            train_task = self.progress.add_task(
                "Training", total=self.config_manager.config.max_iter
            )
            if self.perform_val:
                val_task = self.progress.add_task(
                    "Validation",
                    total=(self.config_manager.config.max_iter + 1) / self.config.val_every,
                )
            else:
                val_task = None
            train_losses, val_losses = self._train_epochs(
                epoch_start, epoch_end, train_task, val_task
            )
    else:
        train_losses, val_losses = self._train_epochs(epoch_start, epoch_end)

    self.on_train_end(train_losses, val_losses)
    return train_losses

`_train_epochs(epoch_start, epoch_end, train_task=None, val_task=None)`

Executes the training loop for a series of epochs.

PARAMETER	DESCRIPTION
`epoch_start`	The starting epoch index. TYPE: `int`
`epoch_end`	The ending epoch index (non-inclusive). TYPE: `int`
`train_task`	The progress bar task ID for training updates. If provided, the progress bar will be updated after each epoch. Defaults to None. TYPE: `int \| None` DEFAULT: `None`
`val_task`	The progress bar task ID for validation updates. If provided and validation is enabled, the progress bar will be updated after each validation run. Defaults to None. TYPE: `int \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`list[list[tuple[Tensor, dict[str, Any]]]]`	list[list[tuple[torch.Tensor, dict[str, Any]]]]: A tuple of
`list[list[tuple[Tensor, dict[str, Any]]]]`	Training loss metrics for all epochs. list -> list -> tuples Epochs -> Training Batches -> (loss, metrics)
`tuple[list[list[tuple[Tensor, dict[str, Any]]]], list[list[tuple[Tensor, dict[str, Any]]]]]`	And Validation loss metrics for all epochs list -> list -> tuples Epochs -> Training Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

def _train_epochs(
    self,
    epoch_start: int,
    epoch_end: int,
    train_task: int | None = None,
    val_task: int | None = None,
) -> tuple[
    list[list[tuple[torch.Tensor, dict[str, Any]]]],
    list[list[tuple[torch.Tensor, dict[str, Any]]]],
]:
    """
    Executes the training loop for a series of epochs.

    Args:
        epoch_start (int): The starting epoch index.
        epoch_end (int): The ending epoch index (non-inclusive).
        train_task (int | None, optional): The progress bar task ID for training updates.
            If provided, the progress bar will be updated after each epoch. Defaults to None.
        val_task (int | None, optional): The progress bar task ID for validation updates.
            If provided and validation is enabled, the progress bar will be updated after each validation run.
            Defaults to None.

    Returns:
        list[list[tuple[torch.Tensor, dict[str, Any]]]]: A tuple of
        Training loss metrics for all epochs.
            list    -> list                  -> tuples
            Epochs  -> Training Batches      -> (loss, metrics)
        And Validation loss metrics for all epochs
            list    -> list                  -> tuples
            Epochs  -> Training Batches      -> (loss, metrics)
    """
    train_losses = []
    val_losses = []

    # Iterate over the epochs
    for epoch in range(epoch_start, epoch_end):
        if not self.stop_training():
            try:
                self.current_epoch = epoch
                self.on_train_epoch_start()
                train_epoch_loss_metrics = self.run_training(self.train_dataloader)
                train_losses.append(train_epoch_loss_metrics)
                self.on_train_epoch_end(train_epoch_loss_metrics)

                # Run validation periodically if specified
                if self.perform_val and (epoch % self.config.val_every == 0):
                    self.on_val_epoch_start()
                    val_epoch_loss_metrics = self.run_validation(self.val_dataloader)
                    val_losses.append(val_epoch_loss_metrics)
                    self.on_val_epoch_end(val_epoch_loss_metrics)
                    if val_task is not None:
                        self.progress.update(val_task, advance=1)  # type: ignore[union-attr]

                if train_task is not None:
                    self.progress.update(train_task, advance=1)  # type: ignore[union-attr]
            except KeyboardInterrupt:
                self._stop_training.fill_(1)
        else:
            if self.accelerator.rank == 0:
                logger.info("Terminating training gracefully after the current iteration.")
            self.accelerator.finalize()
            break
    return train_losses, val_losses

`build_optimize_result(result)`

Builds and stores the optimization result by calculating the average loss and metrics.

Result (or loss_metrics) can have multiple formats: - None Indicates no loss or metrics data is provided. - tuple[torch.Tensor, dict[str, Any]] A single tuple containing the loss tensor and metrics dictionary - at the end of batch. - list[tuple[torch.Tensor, dict[str, Any]]] A list of tuples for multiple batches. - list[list[tuple[torch.Tensor, dict[str, Any]]]] A list of lists of tuples, where each inner list represents metrics across multiple batches within an epoch.

PARAMETER	DESCRIPTION
`result`	(None \| tuple[torch.Tensor, dict[Any, Any]] \| list[tuple[torch.Tensor, dict[Any, Any]]] \| list[list[tuple[torch.Tensor, dict[Any, Any]]]]) The loss and metrics data, which can have multiple formats TYPE: `None \| tuple[Tensor, dict[Any, Any]] \| list[tuple[Tensor, dict[Any, Any]]] \| list[list[tuple[Tensor, dict[Any, Any]]]]`

RETURNS	DESCRIPTION
`None`	This method does not return anything. It sets `self.opt_result` with TYPE: `None`
`None`	the computed average loss and metrics.

Source code in qadence/ml_tools/trainer.py

def build_optimize_result(
    self,
    result: (
        None
        | tuple[torch.Tensor, dict[Any, Any]]
        | list[tuple[torch.Tensor, dict[Any, Any]]]
        | list[list[tuple[torch.Tensor, dict[Any, Any]]]]
    ),
) -> None:
    """
    Builds and stores the optimization result by calculating the average loss and metrics.

    Result (or loss_metrics) can have multiple formats:
    - `None` Indicates no loss or metrics data is provided.
    - `tuple[torch.Tensor, dict[str, Any]]` A single tuple containing the loss tensor
        and metrics dictionary - at the end of batch.
    - `list[tuple[torch.Tensor, dict[str, Any]]]` A list of tuples for
        multiple batches.
    - `list[list[tuple[torch.Tensor, dict[str, Any]]]]` A list of lists of tuples,
    where each inner list represents metrics across multiple batches within an epoch.

    Args:
        result: (None |
                tuple[torch.Tensor, dict[Any, Any]] |
                list[tuple[torch.Tensor, dict[Any, Any]]] |
                list[list[tuple[torch.Tensor, dict[Any, Any]]]])
                    The loss and metrics data, which can have multiple formats

    Returns:
        None: This method does not return anything. It sets `self.opt_result` with
        the computed average loss and metrics.
    """
    loss_metrics = result
    if loss_metrics is None:
        loss = None
        metrics: dict[Any, Any] = {}
    elif isinstance(loss_metrics, tuple):
        # Single tuple case
        loss, metrics = loss_metrics
    else:
        last_epoch: list[tuple[torch.Tensor, dict[Any, Any]]] = []
        if isinstance(loss_metrics, list):
            # Check if it's a list of tuples
            if all(isinstance(item, tuple) for item in loss_metrics):
                last_epoch = cast(list[tuple[torch.Tensor, dict[Any, Any]]], loss_metrics)
            # Check if it's a list of lists of tuples
            elif all(isinstance(item, list) for item in loss_metrics):
                last_epoch = cast(
                    list[tuple[torch.Tensor, dict[Any, Any]]],
                    loss_metrics[-1] if loss_metrics else [],
                )
            else:
                raise ValueError(
                    "Invalid format for result: Expected None, tuple, list of tuples,"
                    " or list of lists of tuples."
                )

        if not last_epoch:
            loss, metrics = None, {}
        else:
            # Compute the average loss over the batches
            loss_tensor = torch.stack([loss_batch for loss_batch, _ in last_epoch])
            avg_loss = loss_tensor.mean()

            # Collect and average metrics for all batches
            metric_keys = last_epoch[0][1].keys()
            metrics_stacked: dict = {key: [] for key in metric_keys}

            for _, metrics_batch in last_epoch:
                for key in metric_keys:
                    value = metrics_batch[key]
                    metrics_stacked[key].append(value)

            avg_metrics = {key: torch.stack(metrics_stacked[key]).mean() for key in metric_keys}

            loss, metrics = avg_loss, avg_metrics

    # Store the optimization result
    self.opt_result = OptimizeResult(
        self.current_epoch,
        self.model,
        self.optimizer,
        loss,
        metrics,
        rank=self.accelerator.rank,
        device=self.accelerator.execution.device,
    )

`fit(train_dataloader=None, val_dataloader=None)`

Fits the model using the specified training configuration.

The dataloaders can be provided to train on new datasets, or the default dataloaders provided in the trainer will be used.

PARAMETER	DESCRIPTION
`train_dataloader`	DataLoader for training data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`val_dataloader`	DataLoader for validation data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`tuple[Module, Optimizer]`	tuple[nn.Module, optim.Optimizer]: The trained model and optimizer.

Source code in qadence/ml_tools/trainer.py

def fit(
    self,
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
) -> tuple[nn.Module, optim.Optimizer]:
    """
    Fits the model using the specified training configuration.

    The dataloaders can be provided to train on new datasets, or the default dataloaders
    provided in the trainer will be used.

    Args:
        train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
        val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.

    Returns:
        tuple[nn.Module, optim.Optimizer]: The trained model and optimizer.
    """
    if train_dataloader is not None:
        self.train_dataloader = train_dataloader
    if val_dataloader is not None:
        self.val_dataloader = val_dataloader

    self._fit_setup()
    self._train()
    self._fit_end()
    self.training_stage = TrainingStage("idle")
    return self.model, self.optimizer

`get_ic_grad_bounds(eta, epsilons, variation_multiple=20, dataloader=None)`

Calculate the bounds on the gradient norm of the loss using Information Content.

PARAMETER	DESCRIPTION
`eta`	The sensitivity IC. TYPE: `float`
`epsilons`	The epsilons to use for thresholds to for discretization of the finite derivatives. TYPE: `Tensor`
`variation_multiple`	The number of sets of variational parameters to generate per each variational parameter. The number of variational parameters required for the statisctiacal analysis scales linearly with the amount of them present in the model. This is that linear factor. TYPE: `int` DEFAULT: `20`
`dataloader`	The dataloader for training data. A new dataloader can be provided, or the dataloader provided in the trinaer will be used. In case no dataloaders are provided at either places, it assumes that the model does not require any input data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`tuple[float, float, float]`	tuple[float, float, float]: The max IC lower bound, max IC upper bound, and sensitivity IC upper bound.

Examples:

import torch
from torch.optim.adam import Adam

from qadence.constructors import ObservableConfig
from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig, TrainConfig
from qadence.ml_tools.data import to_dataloader
from qadence.ml_tools.models import QNN
from qadence.ml_tools.optimize_step import optimize_step
from qadence.ml_tools.trainer import Trainer
from qadence.operations.primitive import Z

fm_config = FeatureMapConfig(num_features=1)
ansatz_config = AnsatzConfig(depth=4)
obs_config = ObservableConfig(detuning=Z)

qnn = QNN.from_configs(
    register=4,
    obs_config=obs_config,
    fm_config=fm_config,
    ansatz_config=ansatz_config,
)

optimizer = Adam(qnn.parameters(), lr=0.001)

batch_size = 25
x = torch.linspace(0, 1, 32).reshape(-1, 1)
y = torch.sin(x)
train_loader = to_dataloader(x, y, batch_size=batch_size, infinite=True)

train_config = TrainConfig(max_iter=100)

trainer = Trainer(
    model=qnn,
    optimizer=optimizer,
    config=train_config,
    loss_fn="mse",
    train_dataloader=train_loader,
    optimize_step=optimize_step,
)

# Perform exploratory landscape analysis with Information Content
ic_sensitivity_threshold = 1e-4
epsilons = torch.logspace(-2, 2, 10)

max_ic_lower_bound, max_ic_upper_bound, sensitivity_ic_upper_bound = (
    trainer.get_ic_grad_bounds(
        eta=ic_sensitivity_threshold,
        epsilons=epsilons,
    )
)

# Resume training as usual...

trainer.fit(train_loader)

Source code in qadence/ml_tools/trainer.py

def get_ic_grad_bounds(
    self,
    eta: float,
    epsilons: torch.Tensor,
    variation_multiple: int = 20,
    dataloader: DataLoader | DictDataLoader | None = None,
) -> tuple[float, float, float]:
    """
    Calculate the bounds on the gradient norm of the loss using Information Content.

    Args:
        eta (float): The sensitivity IC.
        epsilons (torch.Tensor): The epsilons to use for thresholds to for discretization of the
            finite derivatives.
        variation_multiple (int): The number of sets of variational parameters to generate per
            each variational parameter. The number of variational parameters required for the
            statisctiacal analysis scales linearly with the amount of them present in the
            model. This is that linear factor.
        dataloader (DataLoader | DictDataLoader | None): The dataloader for training data. A
            new dataloader can be provided, or the dataloader provided in the trinaer will be
            used. In case no dataloaders are provided at either places, it assumes that the
            model does not require any input data.

    Returns:
        tuple[float, float, float]: The max IC lower bound, max IC upper bound, and sensitivity
            IC upper bound.

    Examples:
        ```python
        import torch
        from torch.optim.adam import Adam

        from qadence.constructors import ObservableConfig
        from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig, TrainConfig
        from qadence.ml_tools.data import to_dataloader
        from qadence.ml_tools.models import QNN
        from qadence.ml_tools.optimize_step import optimize_step
        from qadence.ml_tools.trainer import Trainer
        from qadence.operations.primitive import Z

        fm_config = FeatureMapConfig(num_features=1)
        ansatz_config = AnsatzConfig(depth=4)
        obs_config = ObservableConfig(detuning=Z)

        qnn = QNN.from_configs(
            register=4,
            obs_config=obs_config,
            fm_config=fm_config,
            ansatz_config=ansatz_config,
        )

        optimizer = Adam(qnn.parameters(), lr=0.001)

        batch_size = 25
        x = torch.linspace(0, 1, 32).reshape(-1, 1)
        y = torch.sin(x)
        train_loader = to_dataloader(x, y, batch_size=batch_size, infinite=True)

        train_config = TrainConfig(max_iter=100)

        trainer = Trainer(
            model=qnn,
            optimizer=optimizer,
            config=train_config,
            loss_fn="mse",
            train_dataloader=train_loader,
            optimize_step=optimize_step,
        )

        # Perform exploratory landscape analysis with Information Content
        ic_sensitivity_threshold = 1e-4
        epsilons = torch.logspace(-2, 2, 10)

        max_ic_lower_bound, max_ic_upper_bound, sensitivity_ic_upper_bound = (
            trainer.get_ic_grad_bounds(
                eta=ic_sensitivity_threshold,
                epsilons=epsilons,
            )
        )

        # Resume training as usual...

        trainer.fit(train_loader)
        ```
    """
    if not self._use_grad:
        logger.warning(
            "Gradient norm bounds are only relevant when using a gradient based optimizer. \
                Currently the trainer is set to use a gradient-free optimizer."
        )

    dataloader = dataloader if dataloader is not None else self.train_dataloader

    batch = next(iter(self._batch_iter(dataloader, num_batches=1)))

    ic = InformationContent(self.model, self.loss_fn, batch, epsilons)

    max_ic_lower_bound, max_ic_upper_bound = ic.get_grad_norm_bounds_max_IC()
    sensitivity_ic_upper_bound = ic.get_grad_norm_bounds_sensitivity_IC(eta)

    return max_ic_lower_bound, max_ic_upper_bound, sensitivity_ic_upper_bound

`run_test_batch(batch)`

Runs a single test batch.

PARAMETER	DESCRIPTION
`batch`	Batch of data from the DataLoader. TYPE: `tuple[Tensor, ...]`

RETURNS	DESCRIPTION
`tuple[Tensor, dict[str, Any]]`	tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("test_batch")
def run_test_batch(
    self, batch: tuple[torch.Tensor, ...]
) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Runs a single test batch.

    Args:
        batch (tuple[torch.Tensor, ...]): Batch of data from the DataLoader.

    Returns:
        tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.
    """
    with torch.no_grad():
        loss_metrics = self.loss_fn(self.model, batch)
    return self._modify_batch_end_loss_metrics(loss_metrics)

`run_train_batch(batch)`

Runs a single training batch, performing optimization.

We use the step function to optimize the model based on use_grad. use_grad = True entails gradient based optimization, for which we use optimize_step function. use_grad = False entails gradient free optimization, for which we use update_ng_parameters function.

PARAMETER	DESCRIPTION
`batch`	Batch of data from the DataLoader. TYPE: `tuple[Tensor, ...]`

RETURNS	DESCRIPTION
`tuple[Tensor, dict[str, Any]]`	tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch. tuple of (loss, metrics)

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("train_batch")
def run_train_batch(
    self, batch: tuple[torch.Tensor, ...]
) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Runs a single training batch, performing optimization.

    We use the step function to optimize the model based on use_grad.
        use_grad = True entails gradient based optimization, for which we use
        optimize_step function.
        use_grad = False entails gradient free optimization, for which we use
        update_ng_parameters function.

    Args:
        batch (tuple[torch.Tensor, ...]): Batch of data from the DataLoader.

    Returns:
        tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.
            tuple of (loss, metrics)
    """

    if self.use_grad:
        # Perform gradient-based optimization
        loss_metrics = self.optimize_step(
            model=self.model,
            optimizer=self.optimizer,
            loss_fn=self.loss_fn,
            xs=batch,
            device=self.accelerator.execution.device,
            dtype=self.accelerator.execution.data_dtype,
        )
    else:
        # Perform optimization using Nevergrad
        loss, metrics, ng_params = update_ng_parameters(
            model=self.model,
            optimizer=self.optimizer,
            loss_fn=self.loss_fn,
            data=batch,
            ng_params=self.ng_params,  # type: ignore[arg-type]
        )
        self.ng_params = ng_params
        loss_metrics = loss, metrics

    return self._modify_batch_end_loss_metrics(loss_metrics)

`run_training(dataloader)`

Runs the training for a single epoch, iterating over multiple batches.

PARAMETER	DESCRIPTION
`dataloader`	DataLoader for training data. TYPE: `DataLoader`

RETURNS	DESCRIPTION
`list[tuple[Tensor, dict[str, Any]]]`	list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch. list -> tuples Training Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("train_epoch")
def run_training(self, dataloader: DataLoader) -> list[tuple[torch.Tensor, dict[str, Any]]]:
    """
    Runs the training for a single epoch, iterating over multiple batches.

    Args:
        dataloader (DataLoader): DataLoader for training data.

    Returns:
        list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch.
            list                  -> tuples
            Training Batches      -> (loss, metrics)
    """
    self.model.train()
    train_epoch_loss_metrics = []

    for batch in self._batch_iter(dataloader, self.num_training_batches):
        self.on_train_batch_start(batch)
        train_batch_loss_metrics = self.run_train_batch(batch)
        if self.config.all_reduce_metrics:
            train_batch_loss_metrics = self._aggregate_result(train_batch_loss_metrics)
        train_epoch_loss_metrics.append(train_batch_loss_metrics)
        self.on_train_batch_end(train_batch_loss_metrics)

    return train_epoch_loss_metrics

`run_val_batch(batch)`

Runs a single validation batch.

PARAMETER	DESCRIPTION
`batch`	Batch of data from the DataLoader. TYPE: `tuple[Tensor, ...]`

RETURNS	DESCRIPTION
`tuple[Tensor, dict[str, Any]]`	tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("val_batch")
def run_val_batch(self, batch: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Runs a single validation batch.

    Args:
        batch (tuple[torch.Tensor, ...]): Batch of data from the DataLoader.

    Returns:
        tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.
    """
    with torch.no_grad():
        loss_metrics = self.loss_fn(self.model, batch)
    return self._modify_batch_end_loss_metrics(loss_metrics)

`run_validation(dataloader)`

Runs the validation loop for a single epoch, iterating over multiple batches.

PARAMETER	DESCRIPTION
`dataloader`	DataLoader for validation data. TYPE: `DataLoader`

RETURNS	DESCRIPTION
`list[tuple[Tensor, dict[str, Any]]]`	list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch. list -> tuples Validation Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("val_epoch")
def run_validation(self, dataloader: DataLoader) -> list[tuple[torch.Tensor, dict[str, Any]]]:
    """
    Runs the validation loop for a single epoch, iterating over multiple batches.

    Args:
        dataloader (DataLoader): DataLoader for validation data.

    Returns:
        list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch.
            list                  -> tuples
            Validation Batches      -> (loss, metrics)
    """
    self.model.eval()
    val_epoch_loss_metrics = []

    for batch in self._batch_iter(dataloader, self.num_validation_batches):
        self.on_val_batch_start(batch)
        val_batch_loss_metrics = self.run_val_batch(batch)
        if self.config.all_reduce_metrics:
            val_batch_loss_metrics = self._aggregate_result(val_batch_loss_metrics)
        val_epoch_loss_metrics.append(val_batch_loss_metrics)
        self.on_val_batch_end(val_batch_loss_metrics)

    return val_epoch_loss_metrics

`stop_training()`

Helper function to indicate if the training should be stopped.

We all_reduce the indicator across all processes to ensure all processes are stopped.

Notes

self._stop_training indicator indicates if the training should be stopped. 0 is continue. 1 is stop.

Source code in qadence/ml_tools/trainer.py

def stop_training(self) -> bool:
    """
    Helper function to indicate if the training should be stopped.

    We all_reduce the indicator across all processes to ensure all processes are stopped.

    Notes:
        self._stop_training indicator indicates if the training should be stopped.
        0 is continue. 1 is stop.
    """
    _stop_training = self.accelerator.all_reduce_dict(
        {"indicator": self._stop_training}, op="max"
    )
    return bool(_stop_training["indicator"] > 0)

`test(test_dataloader=None)`

Runs the testing loop if a test DataLoader is provided.

if the test_dataloader is not provided, default test_dataloader defined in the Trainer class is used.

PARAMETER	DESCRIPTION
`test_dataloader`	DataLoader for test data. TYPE: `DataLoader` DEFAULT: `None`

RETURNS	DESCRIPTION
`list[tuple[Tensor, dict[str, Any]]]`	list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch. list -> tuples Test Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

def test(self, test_dataloader: DataLoader = None) -> list[tuple[torch.Tensor, dict[str, Any]]]:
    """
    Runs the testing loop if a test DataLoader is provided.

    if the test_dataloader is not provided, default test_dataloader defined
    in the Trainer class is used.

    Args:
        test_dataloader (DataLoader): DataLoader for test data.

    Returns:
        list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch.
            list                    -> tuples
            Test Batches            -> (loss, metrics)
    """
    if test_dataloader is not None:
        self.test_dataloader = test_dataloader

    self.model.eval()
    test_loss_metrics = []

    for batch in self._batch_iter(test_dataloader, self.num_training_batches):
        self.on_test_batch_start(batch)
        loss_metrics = self.run_test_batch(batch)
        test_loss_metrics.append(loss_metrics)
        self.on_test_batch_end(loss_metrics)

    return test_loss_metrics

`AnsatzConfig(depth=1, ansatz_type=AnsatzType.HEA, ansatz_strategy=Strategy.DIGITAL, strategy_args=dict(), m_block_qubits=None, param_prefix='theta', tag=None)` `dataclass`

`ansatz_strategy = Strategy.DIGITAL` `class-attribute` `instance-attribute`

Ansatz strategy.

Strategy.DIGITAL for fully digital ansatz. Required if ansatz_type is AnsatzType.ALA. Strategy.SDAQC for analog entangling block. Only available for AnsatzType.HEA or AnsatzType.ALA. Strategy.RYDBERG for fully rydberg hea ansatz. Only available for AnsatzType.HEA.

`ansatz_type = AnsatzType.HEA` `class-attribute` `instance-attribute`

What type of ansatz.

AnsatzType.HEA for Hardware Efficient Ansatz. AnsatzType.IIA for Identity Intialized Ansatz. AnsatzType.ALA for Alternating Layer Ansatz.

`depth = 1` `class-attribute` `instance-attribute`

Number of layers of the ansatz.

`m_block_qubits = None` `class-attribute` `instance-attribute`

The number of qubits in the local entangling block of an Alternating Layer Ansatz (ALA).

Only used when ansatz_type is AnsatzType.ALA.

`param_prefix = 'theta'` `class-attribute` `instance-attribute`

The base bame of the variational parameter.

`strategy_args = field(default_factory=dict)` `class-attribute` `instance-attribute`

A dictionary containing keyword arguments to the function creating the ansatz.

Details about each below.

For Strategy.DIGITAL strategy, accepts the following: periodic (bool): if the qubits should be linked periodically. periodic=False is not supported in emu-c. operations (list): list of operations to cycle through in the digital single-qubit rotations of each layer. Defaults to [RX, RY, RX] for hea and [RX, RY] for iia. entangler (AbstractBlock): 2-qubit entangling operation. Supports CNOT, CZ, CRX, CRY, CRZ, CPHASE. Controlld rotations will have variational parameters on the rotation angles. Defaults to CNOT

For Strategy.SDAQC strategy, accepts the following: operations (list): list of operations to cycle through in the digital single-qubit rotations of each layer. Defaults to [RX, RY, RX] for hea and [RX, RY] for iia. entangler (AbstractBlock): Hamiltonian generator for the analog entangling layer. Time parameter is considered variational. Defaults to NN interaction.

For Strategy.RYDBERG strategy, accepts the following: addressable_detuning: whether to turn on the trainable semi-local addressing pattern on the detuning (n_i terms in the Hamiltonian). Defaults to True. addressable_drive: whether to turn on the trainable semi-local addressing pattern on the drive (sigma_i^x terms in the Hamiltonian). Defaults to False. tunable_phase: whether to have a tunable phase to get both sigma^x and sigma^y rotations in the drive term. If False, only a sigma^x term will be included in the drive part of the Hamiltonian generator. Defaults to False.

`tag = None` `class-attribute` `instance-attribute`

String to indicate the name tag of the ansatz.

Defaults to None, in which case no tag will be applied.

`FeatureMapConfig(num_features=0, basis_set=BasisSet.FOURIER, reupload_scaling=ReuploadScaling.CONSTANT, feature_range=None, target_range=None, multivariate_strategy=MultivariateStrategy.PARALLEL, feature_map_strategy=Strategy.DIGITAL, param_prefix=None, num_repeats=0, operation=None, inputs=None, tag=None)` `dataclass`

`basis_set = BasisSet.FOURIER` `class-attribute` `instance-attribute`

Basis set for feature encoding.

Takes qadence.BasisSet. Give a single BasisSet to use the same for all features. Give a dict of (str, BasisSet) where the key is the name of the variable and the value is the BasisSet to use for encoding that feature. BasisSet.FOURIER for Fourier encoding. BasisSet.CHEBYSHEV for Chebyshev encoding.

`feature_map_strategy = Strategy.DIGITAL` `class-attribute` `instance-attribute`

Strategy for feature map.

Accepts DIGITAL, ANALOG or RYDBERG. Defaults to DIGITAL. If the strategy is incompatible with the operation chosen, then operation gets preference and the given strategy is ignored.

`feature_range = None` `class-attribute` `instance-attribute`

Range of data that the input data is assumed to come from.

Give a single tuple to use the same range for all features. Give a dict of (str, tuple) where the key is the name of the variable and the value is the feature range to use for that feature.

`inputs = None` `class-attribute` `instance-attribute`

List that indicates the order of variables of the tensors that are passed.

Optional if a single feature is being encoded, required otherwise. Given input tensors xs = torch.rand(batch_size, input_size:=2) a QNN with inputs=["t", "x"] will assign t, x = xs[:,0], xs[:,1].

`multivariate_strategy = MultivariateStrategy.PARALLEL` `class-attribute` `instance-attribute`

The encoding strategy in case of multi-variate function.

Takes qadence.MultivariateStrategy. If PARALLEL, the features are encoded in one block of rotation gates with the register being split in sub-registers for each feature. If SERIES, the features are encoded sequentially using the full register for each feature, with an ansatz block between them. PARALLEL is allowed only for DIGITAL feature_map_strategy.

`num_features = 0` `class-attribute` `instance-attribute`

Number of feature parameters to be encoded.

Defaults to 0. Thus, no feature parameters are encoded.

`num_repeats = 0` `class-attribute` `instance-attribute`

Number of feature map layers repeated in the data reuploading step.

If all features are to be repeated the same number of times, then can give a single int. For different number of repetitions for each feature, provide a dict of (str, int) where the key is the name of the variable and the value is the number of repetitions for that feature. This amounts to the number of additional reuploads. So if num_repeats is N, the data gets uploaded N+1 times. Defaults to no repetition.

`operation = None` `class-attribute` `instance-attribute`

Type of operation.

Choose among the analog or digital rotations or a custom callable function returning an AnalogBlock instance. If the type of operation is incompatible with the strategy chosen, then operation gets preference and the given strategy is ignored.

`param_prefix = None` `class-attribute` `instance-attribute`

String prefix to create trainable parameters in Feature Map.

A string prefix to create trainable parameters multiplying the feature parameter inside the feature-encoding function. Note that currently this does not take into account the domain of the feature-encoding function. Defaults to None and thus, the feature map is not trainable. Note that this is separate from the name of the parameter. The user can provide a single prefix for all features, and it will be appended by appropriate feature name automatically.

`reupload_scaling = ReuploadScaling.CONSTANT` `class-attribute` `instance-attribute`

Scaling for encoding the same feature on different qubits.

Scaling used to encode the same feature on different qubits in the same layer of the feature maps. Takes qadence.ReuploadScaling. Give a single ReuploadScaling to use the same for all features. Give a dict of (str, ReuploadScaling) where the key is the name of the variable and the value is the ReuploadScaling to use for encoding that feature. ReuploadScaling.CONSTANT for constant scaling. ReuploadScaling.TOWER for linearly increasing scaling. ReuploadScaling.EXP for exponentially increasing scaling.

`tag = None` `class-attribute` `instance-attribute`

String to indicate the name tag of the feature map.

Defaults to None, in which case no tag will be applied.

`target_range = None` `class-attribute` `instance-attribute`

Range of data the data encoder assumes as natural range.

Give a single tuple to use the same range for all features. Give a dict of (str, tuple) where the key is the name of the variable and the value is the target range to use for that feature.

TrainConfig(max_iter=10000, print_every=0, write_every=0, checkpoint_every=0, plot_every=0, callbacks=lambda: list()(), log_model=False, root_folder=Path('./qml_logs'), create_subfolder_per_run=False, log_folder=Path('./'), checkpoint_best_only=False, val_every=0, val_epsilon=1e-05, validation_criterion=None, trainstop_criterion=None, batch_size=1, verbose=True, tracking_tool=ExperimentTrackingTool.TENSORBOARD, hyperparams=dict(), plotting_functions=tuple(), _subfolders=list(), nprocs=1, compute_setup='cpu', backend='gloo', log_setup='cpu', dtype=None, all_reduce_metrics=False) `dataclass`

Default configuration for the training process.

This class provides default settings for various aspects of the training loop, such as logging, checkpointing, and validation. The default values for these fields can be customized when an instance of TrainConfig is created.

Example:

from qadence.ml_tools import TrainConfig
c = TrainConfig(root_folder="/tmp/train")

TrainConfig(max_iter=10000, print_every=0, write_every=0, checkpoint_every=0, plot_every=0, callbacks=[], log_model=False, root_folder='/tmp/train', create_subfolder_per_run=False, log_folder=PosixPath('.'), checkpoint_best_only=False, val_every=0, val_epsilon=1e-05, validation_criterion=None, trainstop_criterion=None, batch_size=1, verbose=True, tracking_tool=<ExperimentTrackingTool.TENSORBOARD: 'tensorboard'>, hyperparams={}, plotting_functions=(), _subfolders=[], nprocs=1, compute_setup='cpu', backend='gloo', log_setup='cpu', dtype=None, all_reduce_metrics=False)

`_subfolders = field(default_factory=list)` `class-attribute` `instance-attribute`

List of subfolders used for logging different runs using the same config inside the.

root folder.

Each subfolder is of structure <id>_<timestamp>_<PID>.

`all_reduce_metrics = False` `class-attribute` `instance-attribute`

Whether to aggregate metrics (e.g., loss, accuracy) across processes.

When True, metrics from different training processes are averaged to provide a consolidated metrics. Note: Since aggregation requires synchronization/all_reduce operation, this can increase the computation time significantly.

`backend = 'gloo'` `class-attribute` `instance-attribute`

Backend used for distributed training communication.

The default is "gloo". Other options may include "nccl" - which is optimized for GPU-based training or "mpi", depending on your system and requirements. It should be one of the backends supported by torch.distributed. For further details, please look at torch backends

`batch_size = 1` `class-attribute` `instance-attribute`

The batch size to use when processing a list or tuple of torch.Tensors.

This specifies how many samples are processed in each training iteration.

`callbacks = field(default_factory=lambda: list())` `class-attribute` `instance-attribute`

List of callbacks to execute during training.

Callbacks can be used for custom behaviors, such as early stopping, custom logging, or other actions triggered at specific events.

`checkpoint_best_only = False` `class-attribute` `instance-attribute`

If True, checkpoints are only saved if there is an improvement in the.

validation metric. This conserves storage by only keeping the best models.

validation_criterion is required when this is set to True.

`checkpoint_every = 0` `class-attribute` `instance-attribute`

Frequency (in epochs) for saving model and optimizer checkpoints during training.

Set to 0 to disable checkpointing. This helps in resuming training or recovering models. Note that setting checkpoint_best_only = True will disable this and only best checkpoints will be saved.

`compute_setup = 'cpu'` `class-attribute` `instance-attribute`

Compute device setup; options are "auto", "gpu", or "cpu".

"auto": Automatically uses GPU if available; otherwise, falls back to CPU.
"gpu": Forces GPU usage, raising an error if no CUDA device is available.
"cpu": Forces the use of CPU regardless of GPU availability.

`create_subfolder_per_run = False` `class-attribute` `instance-attribute`

Whether to create a subfolder for each run, named <id>_<timestamp>_<PID>.

This ensures logs and checkpoints from different runs do not overwrite each other, which is helpful for rapid prototyping. If False, training will resume from the latest checkpoint if one exists in the specified log folder.

`dtype = None` `class-attribute` `instance-attribute`

Data type (precision) for computations.

Both model parameters, and dataset will be of the provided precision.

If not specified or None, the default torch precision (usually torch.float32) is used. If provided dtype is torch.complex128, model parameters will be torch.complex128, and data parameters will be torch.float64

`hyperparams = field(default_factory=dict)` `class-attribute` `instance-attribute`

A dictionary of hyperparameters to be tracked.

This can include learning rates, regularization parameters, or any other training-related configurations.

`log_folder = Path('./')` `class-attribute` `instance-attribute`

The log folder for saving checkpoints and tensorboard logs.

This stores the path where all logs and checkpoints are being saved for this training session. log_folder takes precedence over root_folder, but it is ignored if create_subfolders_per_run=True (in which case, subfolders will be spawned in the root folder).

`log_model = False` `class-attribute` `instance-attribute`

Whether to log a serialized version of the model.

When set to True, the model's state will be logged, useful for model versioning and reproducibility.

`log_setup = 'cpu'` `class-attribute` `instance-attribute`

Logging device setup; options are "auto" or "cpu".

"auto": Uses the same device for logging as for computation.
"cpu": Forces logging to occur on the CPU. This can be useful to avoid potential conflicts with GPU processes.

`max_iter = 10000` `class-attribute` `instance-attribute`

Number of training iterations (epochs) to perform.

This defines the total number of times the model will be updated.

In case of InfiniteTensorDataset, each epoch will have 1 batch. In case of TensorDataset, each epoch will have len(dataloader) batches.

`nprocs = 1` `class-attribute` `instance-attribute`

The number of processes to use for training when spawning subprocesses.

For effective parallel processing, set this to a value greater than 1. - In case of Multi-GPU or Multi-Node-Multi-GPU setups, nprocs should be equal to the total number of GPUs across all nodes (world size), or total number of GPU to be used.

If nprocs > 1, multiple processes will be spawned for training. The training framework will launch additional processes (e.g., for distributed or parallel training). - For CPU setup, this will launch a true parallel processes - For GPU setup, this will launch a distributed training routine. This uses the DistributedDataParallel framework from PyTorch.

`plot_every = 0` `class-attribute` `instance-attribute`

Frequency (in epochs) for generating and saving figures during training.

Set to 0 to disable plotting.

`plotting_functions = field(default_factory=tuple)` `class-attribute` `instance-attribute`

Functions used for in-training plotting.

These are called to generate plots that are logged or saved at specified intervals.

`print_every = 0` `class-attribute` `instance-attribute`

Frequency (in epochs) for printing loss and metrics to the console during training.

Set to 0 to disable this output, meaning that metrics and loss will not be printed during training.

`root_folder = Path('./qml_logs')` `class-attribute` `instance-attribute`

The root folder for saving checkpoints and tensorboard logs.

The default path is "./qml_logs"

This can be set to a specific directory where training artifacts are to be stored. Checkpoints will be saved inside a subfolder in this directory. Subfolders will be created based on create_subfolder_per_run argument.

`tracking_tool = ExperimentTrackingTool.TENSORBOARD` `class-attribute` `instance-attribute`

The tool used for tracking training progress and logging metrics.

Options include tools like TensorBoard, which help visualize and monitor model training.

`trainstop_criterion = None` `class-attribute` `instance-attribute`

A function to determine if the training process should stop based on a.

specific stopping metric. If None, training continues until max_iter is reached.

`val_epsilon = 1e-05` `class-attribute` `instance-attribute`

A small safety margin used to compare the current validation loss with the.

best previous validation loss. This is used to determine improvements in metrics.

`val_every = 0` `class-attribute` `instance-attribute`

Frequency (in epochs) for performing validation.

If set to 0, validation is not performed. Note that metrics from validation are always written, regardless of the write_every setting. Note that initial validation happens at the start of training (when val_every > 0) For initial validation - initial metrics are written. - checkpoint is saved (when checkpoint_best_only = False)

`validation_criterion = None` `class-attribute` `instance-attribute`

A function to evaluate whether a given validation metric meets a desired condition.

The validation_criterion has the following format: def validation_criterion(val_loss: float, best_val_loss: float, val_epsilon: float) -> bool: # process

If None, no custom validation criterion is applied.

`verbose = True` `class-attribute` `instance-attribute`

Whether to print metrics and status messages during training.

If True, detailed metrics and status updates will be displayed in the console.

`write_every = 0` `class-attribute` `instance-attribute`

Frequency (in epochs) for writing loss and metrics using the tracking tool during training.

Set to 0 to disable this logging, which prevents metrics from being logged to the tracking tool. Note that the metrics will always be written at the end of training regardless of this setting.

`get_parameters(model)`

Retrieve all trainable model parameters in a single vector.

PARAMETER	DESCRIPTION
`model`	the input PyTorch model TYPE: `Module`

RETURNS	DESCRIPTION
`Tensor`	a 1-dimensional tensor with the parameters TYPE: `Tensor`

Source code in qadence/ml_tools/parameters.py

def get_parameters(model: Module) -> Tensor:
    """Retrieve all trainable model parameters in a single vector.

    Args:
        model (Module): the input PyTorch model

    Returns:
        Tensor: a 1-dimensional tensor with the parameters
    """
    ps = [p.reshape(-1) for p in model.parameters() if p.requires_grad]
    return torch.concat(ps)

`num_parameters(model)`

Return the total number of parameters of the given model.

Source code in qadence/ml_tools/parameters.py

def num_parameters(model: Module) -> int:
    """Return the total number of parameters of the given model."""
    return len(get_parameters(model))

`set_parameters(model, theta)`

Set all trainable parameters of a model from a single vector.

Notice that this function assumes prior knowledge of right number of parameters in the model

PARAMETER	DESCRIPTION
`model`	the input PyTorch model TYPE: `Module`
`theta`	the parameters to assign TYPE: `Tensor`

Source code in qadence/ml_tools/parameters.py

def set_parameters(model: Module, theta: Tensor) -> None:
    """Set all trainable parameters of a model from a single vector.

    Notice that this function assumes prior knowledge of right number
    of parameters in the model

    Args:
        model (Module): the input PyTorch model
        theta (Tensor): the parameters to assign
    """

    with torch.no_grad():
        idx = 0
        for ps in model.parameters():
            if ps.requires_grad:
                n = torch.numel(ps)
                if ps.ndim == 0:
                    ps[()] = theta[idx : idx + n]
                else:
                    ps[:] = theta[idx : idx + n].reshape(ps.size())
                idx += n

`optimize_step(model, optimizer, loss_fn, xs, device=None, dtype=None)`

Default Torch optimize step with closure.

This is the default optimization step.

PARAMETER	DESCRIPTION
`model`	The input model to be optimized. TYPE: `Module`
`optimizer`	The chosen Torch optimizer. TYPE: `Optimizer`
`loss_fn`	A custom loss function that returns the loss value and a dictionary of metrics. TYPE: `Callable`
`xs`	The input data. If None, it means the given model does not require any input data. TYPE: `dict \| list \| Tensor \| None`
`device`	A target device to run computations on. TYPE: `device` DEFAULT: `None`
`dtype`	Data type for `xs` conversion. TYPE: `dtype` DEFAULT: `None`

RETURNS	DESCRIPTION
`tuple[Tensor \| float, dict \| None]`	tuple[Tensor \| float, dict \| None]: A tuple containing the computed loss value and a dictionary with collected metrics.

Source code in qadence/ml_tools/optimize_step.py

def optimize_step(
    model: Module,
    optimizer: Optimizer,
    loss_fn: Callable,
    xs: dict | list | torch.Tensor | None,
    device: torch.device = None,
    dtype: torch.dtype = None,
) -> tuple[torch.Tensor | float, dict | None]:
    """Default Torch optimize step with closure.

    This is the default optimization step.

    Args:
        model (Module): The input model to be optimized.
        optimizer (Optimizer): The chosen Torch optimizer.
        loss_fn (Callable): A custom loss function
            that returns the loss value and a dictionary of metrics.
        xs (dict | list | Tensor | None): The input data. If None, it means
            the given model does not require any input data.
        device (torch.device): A target device to run computations on.
        dtype (torch.dtype): Data type for `xs` conversion.

    Returns:
        tuple[Tensor | float, dict | None]: A tuple containing the computed loss value
            and a dictionary with collected metrics.
    """

    loss, metrics = None, {}

    def closure() -> Any:
        # NOTE: We need the nonlocal as we can't return a metric dict and
        # because e.g. LBFGS calls this closure multiple times but for some
        # reason the returned loss is always the first one...
        nonlocal metrics, loss
        optimizer.zero_grad()
        loss, metrics = loss_fn(model, xs)
        loss.backward(retain_graph=True)
        return loss.item()

    optimizer.step(closure)
    # return the loss/metrics that are being mutated inside the closure...
    return loss, metrics

`update_ng_parameters(model, optimizer, loss_fn, data, ng_params)`

Update the model parameters using Nevergrad.

This function integrates Nevergrad for derivative-free optimization.

PARAMETER	DESCRIPTION
`model`	The PyTorch model to be optimized. TYPE: `Module`
`optimizer`	A Nevergrad optimizer instance. TYPE: `Optimizer`
`loss_fn`	A custom loss function that returns the loss value and a dictionary of metrics. TYPE: `Callable[[Module, Tensor \| None], tuple[float, dict]]`
`data`	Input data for the model. If None, it means the model does not require input data. TYPE: `Tensor \| None`
`ng_params`	The current set of parameters managed by Nevergrad. TYPE: `Array`

RETURNS	DESCRIPTION
`tuple[float, dict, Array]`	tuple[float, dict, ng.p.Array]: A tuple containing the computed loss value, a dictionary of metrics, and the updated Nevergrad parameters.

Source code in qadence/ml_tools/optimize_step.py

def update_ng_parameters(
    model: Module,
    optimizer: ng.optimizers.Optimizer,
    loss_fn: Callable[[Module, torch.Tensor | None], tuple[float, dict]],
    data: torch.Tensor | None,
    ng_params: ng.p.Array,
) -> tuple[float, dict, ng.p.Array]:
    """Update the model parameters using Nevergrad.

    This function integrates Nevergrad for derivative-free optimization.

    Args:
        model (Module): The PyTorch model to be optimized.
        optimizer (ng.optimizers.Optimizer): A Nevergrad optimizer instance.
        loss_fn (Callable[[Module, Tensor | None], tuple[float, dict]]): A custom loss function
            that returns the loss value and a dictionary of metrics.
        data (Tensor | None): Input data for the model. If None, it means the model does
            not require input data.
        ng_params (ng.p.Array): The current set of parameters managed by Nevergrad.

    Returns:
        tuple[float, dict, ng.p.Array]: A tuple containing the computed loss value,
            a dictionary of metrics, and the updated Nevergrad parameters.
    """
    loss, metrics = loss_fn(model, data)  # type: ignore[misc]
    optimizer.tell(ng_params, float(loss))
    ng_params = optimizer.ask()  # type: ignore[assignment]
    params = promote_to_tensor(ng_params.value, requires_grad=False)
    set_parameters(model, params)
    return loss, metrics, ng_params

`DictDataLoader(dataloaders)` `dataclass`

This class only holds a dictionary of DataLoaders and samples from them.

`InfiniteTensorDataset(*tensors)`

Bases: IterableDataset

Randomly sample points from the first dimension of the given tensors.

Behaves like a normal torch Dataset just that we can sample from it as many times as we want.

Examples:

import torch
from qadence.ml_tools.data import InfiniteTensorDataset

x_data, y_data = torch.rand(5,2), torch.ones(5,1)
# The dataset accepts any number of tensors with the same batch dimension
ds = InfiniteTensorDataset(x_data, y_data)

# call `next` to get one sample from each tensor:
xs = next(iter(ds))

(tensor([0.2952, 0.4321]), tensor([1.]))

Source code in qadence/ml_tools/data.py

def __init__(self, *tensors: Tensor):
    """Randomly sample points from the first dimension of the given tensors.

    Behaves like a normal torch `Dataset` just that we can sample from it as
    many times as we want.

    Examples:
    ```python exec="on" source="above" result="json"
    import torch
    from qadence.ml_tools.data import InfiniteTensorDataset

    x_data, y_data = torch.rand(5,2), torch.ones(5,1)
    # The dataset accepts any number of tensors with the same batch dimension
    ds = InfiniteTensorDataset(x_data, y_data)

    # call `next` to get one sample from each tensor:
    xs = next(iter(ds))
    print(str(xs)) # markdown-exec: hide
    ```
    """
    self.tensors = tensors
    self.indices = list(range(self.tensors[0].size(0)))

`OptimizeResult(iteration, model, optimizer, loss=None, metrics=lambda: dict()(), extra=lambda: dict()(), rank=0, device='cpu')` `dataclass`

OptimizeResult stores many optimization intermediate values.

We store at a current iteration, the model, optimizer, loss values, metrics. An extra dict can be used for saving other information to be used for callbacks.

`device = 'cpu'` `class-attribute` `instance-attribute`

Device on which this result for calculated.

`extra = field(default_factory=lambda: dict())` `class-attribute` `instance-attribute`

Extra dict for saving anything else to be used in callbacks.

`iteration` `instance-attribute`

Current iteration number.

`loss = None` `class-attribute` `instance-attribute`

Loss value.

`metrics = field(default_factory=lambda: dict())` `class-attribute` `instance-attribute`

Metrics that can be saved during training.

`model` `instance-attribute`

Model at iteration.

`optimizer` `instance-attribute`

Optimizer at iteration.

`rank = 0` `class-attribute` `instance-attribute`

Rank of the process for which this result was generated.

`data_to_device(xs, *args, **kwargs)`

Utility method to move arbitrary data to 'device'.

Source code in qadence/ml_tools/data.py

@singledispatch
def data_to_device(xs: Any, *args: Any, **kwargs: Any) -> Any:
    """Utility method to move arbitrary data to 'device'."""
    raise ValueError(f"Unable to move {type(xs)} with input args: {args} and kwargs: {kwargs}.")

`to_dataloader(*tensors, batch_size=1, infinite=False)`

Convert torch tensors an (infinite) Dataloader.

PARAMETER	DESCRIPTION
`*tensors`	Torch tensors to use in the dataloader. TYPE: `Tensor` DEFAULT: `()`
`batch_size`	batch size of sampled tensors TYPE: `int` DEFAULT: `1`
`infinite`	if `True`, the dataloader will keep sampling indefinitely even after the whole dataset was sampled once TYPE: `bool` DEFAULT: `False`

Examples:

import torch
from qadence.ml_tools import to_dataloader

(x, y, z) = [torch.rand(10) for _ in range(3)]
loader = iter(to_dataloader(x, y, z, batch_size=5, infinite=True))
print(next(loader))
print(next(loader))
print(next(loader))

[tensor([0.4483, 0.8226, 0.3765, 0.1782, 0.0435]), tensor([0.8529, 0.4461, 0.4437, 0.3194, 0.8547]), tensor([0.8719, 0.6678, 0.4718, 0.9445, 0.1859])]
[tensor([0.3387, 0.6516, 0.8960, 0.8684, 0.5089]), tensor([0.2403, 0.3733, 0.5831, 0.8063, 0.0280]), tensor([0.2211, 0.4503, 0.2892, 0.3290, 0.9851])]
[tensor([0.4483, 0.8226, 0.3765, 0.1782, 0.0435]), tensor([0.8529, 0.4461, 0.4437, 0.3194, 0.8547]), tensor([0.8719, 0.6678, 0.4718, 0.9445, 0.1859])]

Source code in qadence/ml_tools/data.py

def to_dataloader(*tensors: Tensor, batch_size: int = 1, infinite: bool = False) -> DataLoader:
    """Convert torch tensors an (infinite) Dataloader.

    Arguments:
        *tensors: Torch tensors to use in the dataloader.
        batch_size: batch size of sampled tensors
        infinite: if `True`, the dataloader will keep sampling indefinitely even after the whole
            dataset was sampled once

    Examples:

    ```python exec="on" source="above" result="json"
    import torch
    from qadence.ml_tools import to_dataloader

    (x, y, z) = [torch.rand(10) for _ in range(3)]
    loader = iter(to_dataloader(x, y, z, batch_size=5, infinite=True))
    print(next(loader))
    print(next(loader))
    print(next(loader))
    ```
    """
    ds = InfiniteTensorDataset(*tensors) if infinite else TensorDataset(*tensors)
    return DataLoader(ds, batch_size=batch_size)

`QNN(circuit, observable, backend=BackendName.PYQTORCH, diff_mode=DiffMode.AD, measurement=None, noise=None, configuration=None, inputs=None, input_diff_mode=InputDiffMode.AD)`

Bases: QuantumModel

Quantum neural network model for n-dimensional inputs.

Examples:

import torch
from qadence import QuantumCircuit, QNN, Z
from qadence import hea, feature_map, hamiltonian_factory, kron

# create the circuit
n_qubits, depth = 2, 4
fm = kron(
    feature_map(1, support=(0,), param="x"),
    feature_map(1, support=(1,), param="y")
)
ansatz = hea(n_qubits=n_qubits, depth=depth)
circuit = QuantumCircuit(n_qubits, fm, ansatz)
obs_base = hamiltonian_factory(n_qubits, detuning=Z)

# the QNN will yield two outputs
obs = [2.0 * obs_base, 4.0 * obs_base]

# initialize and use the model
qnn = QNN(circuit, obs, inputs=["x", "y"])
y = qnn(torch.rand(3, 2))

tensor([[0.5935, 1.1870],
        [0.8289, 1.6579],
        [0.4910, 0.9820]], grad_fn=<CatBackward0>)

Initialize the QNN.

The number of inputs is determined by the feature parameters in the input quantum circuit while the number of outputs is determined by how many observables are provided as input

PARAMETER	DESCRIPTION
`circuit`	The quantum circuit to use for the QNN. TYPE: `QuantumCircuit`
`observable`	The observable. TYPE: `list[AbstractBlock] \| AbstractBlock`
`backend`	The chosen quantum backend. TYPE: `BackendName` DEFAULT: `PYQTORCH`
`diff_mode`	The differentiation engine to use. Choices 'gpsr' or 'ad'. TYPE: `DiffMode` DEFAULT: `AD`
`measurement`	optional measurement protocol. If None, use exact expectation value with a statevector simulator TYPE: `Measurements \| None` DEFAULT: `None`
`noise`	A noise model to use. TYPE: `NoiseHandler \| None` DEFAULT: `None`
`configuration`	optional configuration for the backend TYPE: `BackendConfiguration \| dict \| None` DEFAULT: `None`
`inputs`	List that indicates the order of variables of the tensors that are passed to the model. Given input tensors `xs = torch.rand(batch_size, input_size:=2)` a QNN with `inputs=["t", "x"]` will assign `t, x = xs[:,0], xs[:,1]`. TYPE: `list[Basic \| str] \| None` DEFAULT: `None`
`input_diff_mode`	The differentiation mode for the input tensor. TYPE: `InputDiffMode \| str` DEFAULT: `AD`

Source code in qadence/ml_tools/models.py

def __init__(
    self,
    circuit: QuantumCircuit,
    observable: list[AbstractBlock] | AbstractBlock,
    backend: BackendName = BackendName.PYQTORCH,
    diff_mode: DiffMode = DiffMode.AD,
    measurement: Measurements | None = None,
    noise: NoiseHandler | None = None,
    configuration: BackendConfiguration | dict | None = None,
    inputs: list[sympy.Basic | str] | None = None,
    input_diff_mode: InputDiffMode | str = InputDiffMode.AD,
):
    """Initialize the QNN.

    The number of inputs is determined by the feature parameters in the input
    quantum circuit while the number of outputs is determined by how many
    observables are provided as input

    Args:
        circuit: The quantum circuit to use for the QNN.
        observable: The observable.
        backend: The chosen quantum backend.
        diff_mode: The differentiation engine to use. Choices 'gpsr' or 'ad'.
        measurement: optional measurement protocol. If None,
            use exact expectation value with a statevector simulator
        noise: A noise model to use.
        configuration: optional configuration for the backend
        inputs: List that indicates the order of variables of the tensors that are passed
            to the model. Given input tensors `xs = torch.rand(batch_size, input_size:=2)` a QNN
            with `inputs=["t", "x"]` will assign `t, x = xs[:,0], xs[:,1]`.
        input_diff_mode: The differentiation mode for the input tensor.
    """
    super().__init__(
        circuit,
        observable=observable,
        backend=backend,
        diff_mode=diff_mode,
        measurement=measurement,
        configuration=configuration,
        noise=noise,
    )
    if self._observable is None:
        raise ValueError("You need to provide at least one observable in the QNN constructor")
    if (inputs is not None) and (len(self.inputs) == len(inputs)):
        self.inputs = [sympy.symbols(x) if isinstance(x, str) else x for x in inputs]  # type: ignore[union-attr]
    elif (inputs is None) and len(self.inputs) <= 1:
        self.inputs = [sympy.symbols(x) if isinstance(x, str) else x for x in self.inputs]  # type: ignore[union-attr]
    else:
        raise ValueError(
            """
            Your QNN has more than one input. Please provide a list of inputs in the order of
            your tensor domain. For example, if you want to pass
            `xs = torch.rand(batch_size, input_size:=3)` to you QNN, where
            ```
            t = x[:,0]
            x = x[:,1]
            y = x[:,2]
            ```
            you have to specify
            ```
            QNN(circuit, observable, inputs=["t", "x", "y"])
            ```
            You can also pass a list of sympy symbols.
        """
        )
    self.format_to_dict = format_to_dict_fn(self.inputs)  # type: ignore[arg-type]
    self.input_diff_mode = InputDiffMode(input_diff_mode)
    if self.input_diff_mode == InputDiffMode.FD:
        from qadence.backends.utils import finitediff

        self.__derivative = finitediff
    elif self.input_diff_mode == InputDiffMode.AD:
        self.__derivative = _torch_derivative  # type: ignore[assignment]
    else:
        raise ValueError(f"Unkown forward diff mode: {self.input_diff_mode}")

    self._model_configs: dict = dict()

`str()`

Return a string representation of a QNN.

When creating a QNN from a set of configurations, we print the configurations used. Otherwise, we use the default printing.

RETURNS	DESCRIPTION
`str \| Any`	str \| Any: A string representation of a QNN.

Example:

from qadence import QNN
from qadence.constructors.hamiltonians import Interaction
from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig
from qadence.ml_tools.constructors import (
    ObservableConfig,
)
from qadence.operations import Z
from qadence.types import BackendName

backend = BackendName.PYQTORCH
fm_config = FeatureMapConfig(num_features=1)
ansatz_config = AnsatzConfig()
observable_config = ObservableConfig(detuning=Z, interaction=Interaction.ZZ, scale=2)

qnn = QNN.from_configs(
    register=2,
    obs_config=observable_config,
    fm_config=fm_config,
    ansatz_config=ansatz_config,
    backend=backend,
)

QNN(
ansatz_config = AnsatzConfig(depth=1, ansatz_type=<AnsatzType.HEA: 'hea'>, ansatz_strategy=<Strategy.DIGITAL: 'Digital'>, strategy_args={}, m_block_qubits=None, param_prefix='theta', tag=None)
fm_config = FeatureMapConfig(num_features=1, basis_set={'x': <BasisSet.FOURIER: 'Fourier'>}, reupload_scaling={'x': <ReuploadScaling.CONSTANT: 'Constant'>}, feature_range={'x': None}, target_range={'x': None}, multivariate_strategy=<MultivariateStrategy.PARALLEL: 'Parallel'>, feature_map_strategy=<Strategy.DIGITAL: 'Digital'>, param_prefix=None, num_repeats={'x': 0}, operation=<class 'qadence.operations.parametric.RX'>, inputs=['x'], tag=None)
register = 2
observable_config = {'Obs.': '(2.000 * (Z(0) + Z(1) + (Z(0) ⊗ Z(1))))'}
)

Source code in qadence/ml_tools/models.py

def __str__(self) -> str | Any:
    """Return a string representation of a QNN.

    When creating a QNN from a set of configurations,
    we print the configurations used. Otherwise, we use the default printing.

    Returns:
        str | Any: A string representation of a QNN.

    Example:
    ```python exec="on" source="material-block" result="json"
    from qadence import QNN
    from qadence.constructors.hamiltonians import Interaction
    from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig
    from qadence.ml_tools.constructors import (
        ObservableConfig,
    )
    from qadence.operations import Z
    from qadence.types import BackendName

    backend = BackendName.PYQTORCH
    fm_config = FeatureMapConfig(num_features=1)
    ansatz_config = AnsatzConfig()
    observable_config = ObservableConfig(detuning=Z, interaction=Interaction.ZZ, scale=2)

    qnn = QNN.from_configs(
        register=2,
        obs_config=observable_config,
        fm_config=fm_config,
        ansatz_config=ansatz_config,
        backend=backend,
    )
    print(qnn) # markdown-exec: hide
    ```
    """
    if bool(self._model_configs):
        configs_str = "\n".join(
            (
                k + " = " + str(self._model_configs[k])
                for k in sorted(self._model_configs.keys())
                if k != "observable_config"
            )
        )
        observable_str = ""
        if self._observable:
            observable_str = f"observable_config = {self.observables_to_expression()}"

        return f"{type(self).__name__}(\n{configs_str}\n{observable_str}\n)"

    return super().__str__()

`forward(values=None, state=None, measurement=None, noise=None, endianness=Endianness.BIG)`

Forward pass of the model.

This returns the (differentiable) expectation value of the given observable operator defined in the constructor. Differently from the base QuantumModel class, the QNN accepts also a tensor as input for the forward pass. The tensor is expected to have shape: n_batches x in_features where n_batches is the number of data points and in_features is the dimensionality of the problem

The output of the forward pass is the expectation value of the input observable(s). If a single observable is given, the output shape is n_batches while if multiple observables are given the output shape is instead n_batches x n_observables

PARAMETER	DESCRIPTION
`values`	the values of the feature parameters TYPE: `dict[str, Tensor] \| Tensor` DEFAULT: `None`
`state`	Initial state. TYPE: `Tensor \| None` DEFAULT: `None`
`measurement`	optional measurement protocol. If None, use exact expectation value with a statevector simulator TYPE: `Measurements \| None` DEFAULT: `None`
`noise`	A noise model to use. TYPE: `NoiseHandler \| None` DEFAULT: `None`
`endianness`	Endianness of the resulting bit strings. TYPE: `Endianness` DEFAULT: `BIG`

RETURNS	DESCRIPTION
`Tensor`	a tensor with the expectation value of the observables passed in the constructor of the model TYPE: `Tensor`

Source code in qadence/ml_tools/models.py

def forward(
    self,
    values: dict[str, Tensor] | Tensor = None,
    state: Tensor | None = None,
    measurement: Measurements | None = None,
    noise: NoiseHandler | None = None,
    endianness: Endianness = Endianness.BIG,
) -> Tensor:
    """Forward pass of the model.

    This returns the (differentiable) expectation value of the given observable
    operator defined in the constructor. Differently from the base QuantumModel
    class, the QNN accepts also a tensor as input for the forward pass. The
    tensor is expected to have shape: `n_batches x in_features` where `n_batches`
    is the number of data points and `in_features` is the dimensionality of the problem

    The output of the forward pass is the expectation value of the input
    observable(s). If a single observable is given, the output shape is
    `n_batches` while if multiple observables are given the output shape
    is instead `n_batches x n_observables`

    Args:
        values: the values of the feature parameters
        state: Initial state.
        measurement: optional measurement protocol. If None,
            use exact expectation value with a statevector simulator
        noise: A noise model to use.
        endianness: Endianness of the resulting bit strings.

    Returns:
        Tensor: a tensor with the expectation value of the observables passed
            in the constructor of the model
    """
    return self.expectation(
        values, state=state, measurement=measurement, noise=noise, endianness=endianness
    )

`from_configs(register, obs_config, fm_config=FeatureMapConfig(), ansatz_config=AnsatzConfig(), backend=BackendName.PYQTORCH, diff_mode=DiffMode.AD, measurement=None, noise=None, configuration=None, input_diff_mode=InputDiffMode.AD)` `classmethod`

Create a QNN from a set of configurations.

PARAMETER	DESCRIPTION
`register`	The number of qubits or a register object. TYPE: `int \| Register`
`obs_config`	The configuration(s) for the observable(s). TYPE: `list[ObservableConfig] \| ObservableConfig`
`fm_config`	The configuration for the feature map. Defaults to no feature encoding block. TYPE: `FeatureMapConfig` DEFAULT: `FeatureMapConfig()`
`ansatz_config`	The configuration for the ansatz. Defaults to a single layer of hardware efficient ansatz. TYPE: `AnsatzConfig` DEFAULT: `AnsatzConfig()`
`backend`	The chosen quantum backend. TYPE: `BackendName` DEFAULT: `PYQTORCH`
`diff_mode`	The differentiation engine to use. Choices are 'gpsr' or 'ad'. TYPE: `DiffMode` DEFAULT: `AD`
`measurement`	Optional measurement protocol. If None, use exact expectation value with a statevector simulator. TYPE: `Measurements` DEFAULT: `None`
`noise`	A noise model to use. TYPE: `Noise` DEFAULT: `None`
`configuration`	Optional backend configuration. TYPE: `BackendConfiguration \| dict` DEFAULT: `None`
`input_diff_mode`	The differentiation mode for the input tensor. TYPE: `InputDiffMode` DEFAULT: `AD`

RETURNS	DESCRIPTION
`QNN`	A QNN object.

RAISES	DESCRIPTION
`ValueError`	If the observable configuration is not provided.

Example:

import torch
from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig
from qadence.ml_tools import QNN
from qadence.constructors import ObservableConfig
from qadence.operations import Z
from qadence.types import (
    AnsatzType, BackendName, BasisSet, ReuploadScaling, Strategy
)

register = 4
obs_config = ObservableConfig(
    detuning=Z,
    scale=5.0,
    shift=0.0,
    trainable_transform=None,
)
fm_config = FeatureMapConfig(
    num_features=2,
    inputs=["x", "y"],
    basis_set=BasisSet.FOURIER,
    reupload_scaling=ReuploadScaling.CONSTANT,
    feature_range={
        "x": (-1.0, 1.0),
        "y": (0.0, 1.0),
    },
)
ansatz_config = AnsatzConfig(
    depth=2,
    ansatz_type=AnsatzType.HEA,
    ansatz_strategy=Strategy.DIGITAL,
)

qnn = QNN.from_configs(
    register, obs_config, fm_config, ansatz_config, backend=BackendName.PYQTORCH
)

x = torch.rand(2, 2)
y = qnn(x)

tensor([[-2.0401],
        [ 4.5571]], grad_fn=<CatBackward0>)

Source code in qadence/ml_tools/models.py

@classmethod
def from_configs(
    cls,
    register: int | Register,
    obs_config: Any,
    fm_config: Any = FeatureMapConfig(),
    ansatz_config: Any = AnsatzConfig(),
    backend: BackendName = BackendName.PYQTORCH,
    diff_mode: DiffMode = DiffMode.AD,
    measurement: Measurements | None = None,
    noise: NoiseHandler | None = None,
    configuration: BackendConfiguration | dict | None = None,
    input_diff_mode: InputDiffMode | str = InputDiffMode.AD,
) -> QNN:
    """Create a QNN from a set of configurations.

    Args:
        register (int | Register): The number of qubits or a register object.
        obs_config (list[ObservableConfig] | ObservableConfig): The configuration(s)
            for the observable(s).
        fm_config (FeatureMapConfig): The configuration for the feature map.
            Defaults to no feature encoding block.
        ansatz_config (AnsatzConfig): The configuration for the ansatz.
            Defaults to a single layer of hardware efficient ansatz.
        backend (BackendName): The chosen quantum backend.
        diff_mode (DiffMode): The differentiation engine to use. Choices are
            'gpsr' or 'ad'.
        measurement (Measurements): Optional measurement protocol. If None,
            use exact expectation value with a statevector simulator.
        noise (Noise): A noise model to use.
        configuration (BackendConfiguration | dict): Optional backend configuration.
        input_diff_mode (InputDiffMode): The differentiation mode for the input tensor.

    Returns:
        A QNN object.

    Raises:
        ValueError: If the observable configuration is not provided.

    Example:
    ```python exec="on" source="material-block" result="json"
    import torch
    from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig
    from qadence.ml_tools import QNN
    from qadence.constructors import ObservableConfig
    from qadence.operations import Z
    from qadence.types import (
        AnsatzType, BackendName, BasisSet, ReuploadScaling, Strategy
    )

    register = 4
    obs_config = ObservableConfig(
        detuning=Z,
        scale=5.0,
        shift=0.0,
        trainable_transform=None,
    )
    fm_config = FeatureMapConfig(
        num_features=2,
        inputs=["x", "y"],
        basis_set=BasisSet.FOURIER,
        reupload_scaling=ReuploadScaling.CONSTANT,
        feature_range={
            "x": (-1.0, 1.0),
            "y": (0.0, 1.0),
        },
    )
    ansatz_config = AnsatzConfig(
        depth=2,
        ansatz_type=AnsatzType.HEA,
        ansatz_strategy=Strategy.DIGITAL,
    )

    qnn = QNN.from_configs(
        register, obs_config, fm_config, ansatz_config, backend=BackendName.PYQTORCH
    )

    x = torch.rand(2, 2)
    y = qnn(x)
    print(str(y)) # markdown-exec: hide
    ```
    """
    from .constructors import build_qnn_from_configs

    qnn = build_qnn_from_configs(
        register=register,
        observable_config=obs_config,
        fm_config=fm_config,
        ansatz_config=ansatz_config,
        backend=backend,
        diff_mode=diff_mode,
        measurement=measurement,
        noise=noise,
        configuration=configuration,
        input_diff_mode=input_diff_mode,
    )
    qnn._model_configs = {
        "register": register,
        "observable_config": obs_config,
        "fm_config": fm_config,
        "ansatz_config": ansatz_config,
    }
    return qnn

`derivative(ufa, x, derivative_indices)`

Compute derivatives w.r.t.

inputs of a UFA with a single output. The derivative_indices specify which derivative(s) are computed. E.g. derivative_indices=(1,2) would compute the a second order derivative w.r.t to the indices 1 and 2 of the input tensor.

PARAMETER	DESCRIPTION
`ufa`	The model for which we want to compute the derivative. TYPE: `Module`
`x`	(batch_size, input_size) input tensor. TYPE: `Tensor`
`derivative_indices`	Define which derivatives to compute. TYPE: `tuple`

Examples: If we create a UFA with three inputs and denote the first, second, and third input with x, y, and z we can compute the following derivatives w.r.t to those inputs:

import torch
from qadence.ml_tools.models import derivative, QNN
from qadence.ml_tools.config import FeatureMapConfig, AnsatzConfig
from qadence.constructors.hamiltonians import ObservableConfig
from qadence.operations import Z

fm_config = FeatureMapConfig(num_features=3, inputs=["x", "y", "z"])
ansatz_config = AnsatzConfig()
obs_config = ObservableConfig(detuning=Z)

f = QNN.from_configs(
    register=3, obs_config=obs_config, fm_config=fm_config, ansatz_config=ansatz_config,
)
inputs = torch.rand(5,3,requires_grad=True)

# df_dx
derivative(f, inputs, (0,))

# d2f_dydz
derivative(f, inputs, (1,2))

# d3fdy2dx
derivative(f, inputs, (1,1,0))

Source code in qadence/ml_tools/models.py

def derivative(ufa: torch.nn.Module, x: Tensor, derivative_indices: tuple[int, ...]) -> Tensor:
    """Compute derivatives w.r.t.

    inputs of a UFA with a single output. The
    `derivative_indices` specify which derivative(s) are computed.  E.g.
    `derivative_indices=(1,2)` would compute the a second order derivative w.r.t
    to the indices `1` and `2` of the input tensor.

    Arguments:
        ufa: The model for which we want to compute the derivative.
        x (Tensor): (batch_size, input_size) input tensor.
        derivative_indices (tuple): Define which derivatives to compute.

    Examples:
    If we create a UFA with three inputs and denote the first, second, and third
    input with `x`, `y`, and `z` we can compute the following derivatives w.r.t
    to those inputs:
    ```py exec="on" source="material-block"
    import torch
    from qadence.ml_tools.models import derivative, QNN
    from qadence.ml_tools.config import FeatureMapConfig, AnsatzConfig
    from qadence.constructors.hamiltonians import ObservableConfig
    from qadence.operations import Z

    fm_config = FeatureMapConfig(num_features=3, inputs=["x", "y", "z"])
    ansatz_config = AnsatzConfig()
    obs_config = ObservableConfig(detuning=Z)

    f = QNN.from_configs(
        register=3, obs_config=obs_config, fm_config=fm_config, ansatz_config=ansatz_config,
    )
    inputs = torch.rand(5,3,requires_grad=True)

    # df_dx
    derivative(f, inputs, (0,))

    # d2f_dydz
    derivative(f, inputs, (1,2))

    # d3fdy2dx
    derivative(f, inputs, (1,1,0))
    ```
    """
    assert ufa.out_features == 1, "Can only call `derivative` on models with 1D output."
    return ufa._derivative(x, derivative_indices)

`format_to_dict_fn(inputs=[])`

Format an input tensor into the format required by the forward pass.

The tensor is assumed to have dimensions: n_batches x in_features where in_features corresponds to the number of input features of the QNN

Source code in qadence/ml_tools/models.py

def format_to_dict_fn(
    inputs: list[sympy.Symbol | str] = [],
) -> Callable[[Tensor | ParamDictType], ParamDictType]:
    """Format an input tensor into the format required by the forward pass.

    The tensor is assumed to have dimensions: n_batches x in_features where in_features
    corresponds to the number of input features of the QNN
    """
    in_features = len(inputs)

    def tensor_to_dict(values: Tensor | ParamDictType) -> ParamDictType:
        if isinstance(values, Tensor):
            values = values.reshape(-1, 1) if len(values.size()) == 1 else values
            if not values.shape[1] == in_features:
                raise ValueError(
                    f"Model expects in_features={in_features} but got {values.shape[1]}."
                )
            values = {fparam.name: values[:, inputs.index(fparam)] for fparam in inputs}  # type: ignore[union-attr]
        return values

    return tensor_to_dict

`Callback(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Base class for defining various training callbacks.

ATTRIBUTE	DESCRIPTION
`on`	The event on which to trigger the callback. Must be a valid on value from: ["train_start", "train_end", "train_epoch_start", "train_epoch_end", "train_batch_start", "train_batch_end","val_epoch_start", "val_epoch_end", "val_batch_start", "val_batch_end", "test_batch_start", "test_batch_end"] TYPE: `str`
`called_every`	Frequency of callback calls in terms of iterations. TYPE: `int`
`callback`	The function to call if the condition is met. TYPE: `CallbackFunction \| None`
`callback_condition`	Condition to check before calling. TYPE: `CallbackConditionFunction \| None`
`modify_optimize_result`	Function to modify `OptimizeResult`. TYPE: `CallbackFunction \| dict[str, Any] \| None`

A callback can be defined in two ways:

By providing a callback function directly in the base class: This is useful for simple callbacks that don't require subclassing.

Example:

from qadence.ml_tools.callbacks import Callback

def custom_callback_function(trainer, config, writer):
    print("Custom callback executed.")

custom_callback = Callback(
    on="train_end",
    called_every=5,
    callback=custom_callback_function
)

By inheriting and implementing the run_callback method: This is suitable for more complex callbacks that require customization.

Example:

from qadence.ml_tools.callbacks import Callback
class CustomCallback(Callback):
    def run_callback(self, trainer, config, writer):
        print("Custom behavior in the inherited run_callback method.")

custom_callback = CustomCallback(on="train_end", called_every=10)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`on` `property` `writable`

Returns the TrainingStage.

RETURNS	DESCRIPTION
`TrainingStage`	TrainingStage for the callback TYPE: `TrainingStage \| str`

`call(when, trainer, config, writer)`

Executes the callback if conditions are met.

PARAMETER	DESCRIPTION
`when`	The event when the callback is triggered. TYPE: `str`
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

RETURNS	DESCRIPTION
`Any`	Result of the callback function if executed. TYPE: `Any`

Source code in qadence/ml_tools/callbacks/callback.py

def __call__(
    self, when: TrainingStage, trainer: Any, config: TrainConfig, writer: BaseWriter
) -> Any:
    """Executes the callback if conditions are met.

    Args:
        when (str): The event when the callback is triggered.
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.

    Returns:
        Any: Result of the callback function if executed.
    """
    opt_result = trainer.opt_result
    if self.on == when:
        if opt_result:
            opt_result = self.modify_optimize_result(opt_result)
        if self._should_call(when, opt_result):
            return self.run_callback(trainer, config, writer)

`_should_call(when, opt_result)`

Checks if the callback should be called.

PARAMETER	DESCRIPTION
`when`	The event when the callback is considered for execution. TYPE: `str`
`opt_result`	The current optimization results. TYPE: `OptimizeResult`

RETURNS	DESCRIPTION
`bool`	Whether the callback should be called. TYPE: `bool`

Source code in qadence/ml_tools/callbacks/callback.py

def _should_call(self, when: str, opt_result: OptimizeResult) -> bool:
    """Checks if the callback should be called.

    Args:
        when (str): The event when the callback is considered for execution.
        opt_result (OptimizeResult): The current optimization results.

    Returns:
        bool: Whether the callback should be called.
    """
    if when in [TrainingStage("train_start"), TrainingStage("train_end")]:
        return True
    if self.called_every == 0 or opt_result.iteration == 0:
        return False
    if opt_result.iteration % self.called_every == 0 and self.callback_condition(opt_result):
        return True
    return False

`run_callback(trainer, config, writer)`

Executes the defined callback.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

RETURNS	DESCRIPTION
`Any`	Result of the callback execution. TYPE: `Any`

RAISES	DESCRIPTION
`NotImplementedError`	If not implemented in subclasses.

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Executes the defined callback.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.

    Returns:
        Any: Result of the callback execution.

    Raises:
        NotImplementedError: If not implemented in subclasses.
    """
    if self.callback is not None:
        return self.callback(trainer, config, writer)
    raise NotImplementedError("Subclasses should override the run_callback method.")

`EarlyStopping(on, called_every, monitor, patience=5, mode='min')`

Bases: Callback

Stops training when a monitored metric has not improved for a specified number of epochs.

This callback monitors a specified metric (e.g., validation loss or accuracy). If the metric does not improve for a given patience period, training is stopped.

Example Usage in TrainConfig: To use EarlyStopping, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import EarlyStopping

# Create an instance of the EarlyStopping callback
early_stopping = EarlyStopping(on="val_epoch_end",
                               called_every=1,
                               monitor="val_loss",
                               patience=5,
                               mode="min")

config = TrainConfig(
    max_iter=10000,
    print_every=1000,
    callbacks=[early_stopping]
)

Initializes the EarlyStopping callback.

PARAMETER	DESCRIPTION
`on`	The event to trigger the callback (e.g., "val_epoch_end"). TYPE: `str`
`called_every`	Frequency of callback calls in terms of iterations. TYPE: `int`
`monitor`	The metric to monitor (e.g., "val_loss" or "train_loss"). All metrics returned by optimize step are available to monitor. Please add "val_" and "train_" strings at the start of the metric name. TYPE: `str`
`patience`	Number of iterations to wait for improvement. Default is 5. TYPE: `int` DEFAULT: `5`
`mode`	Whether to minimize ("min") or maximize ("max") the metric. Default is "min". TYPE: `str` DEFAULT: `'min'`

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self, on: str, called_every: int, monitor: str, patience: int = 5, mode: str = "min"
):
    """Initializes the EarlyStopping callback.

    Args:
        on (str): The event to trigger the callback (e.g., "val_epoch_end").
        called_every (int): Frequency of callback calls in terms of iterations.
        monitor (str): The metric to monitor (e.g., "val_loss" or "train_loss").
            All metrics returned by optimize step are available to monitor.
            Please add "val_" and "train_" strings at the start of the metric name.
        patience (int, optional): Number of iterations to wait for improvement. Default is 5.
        mode (str, optional): Whether to minimize ("min") or maximize ("max") the metric.
            Default is "min".
    """
    super().__init__(on=on, called_every=called_every)
    self.monitor = monitor
    self.patience = patience
    self.mode = mode
    self.best_value = float("inf") if mode == "min" else -float("inf")
    self.counter = 0

`run_callback(trainer, config, writer)`

Monitors the metric and stops training if no improvement is observed.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
    """
    Monitors the metric and stops training if no improvement is observed.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter): The writer object for logging.
    """
    current_value = trainer.opt_result.metrics.get(self.monitor)
    if current_value is None:
        raise ValueError(f"Metric '{self.monitor}' is not available in the trainer's metrics.")

    if (self.mode == "min" and current_value < self.best_value) or (
        self.mode == "max" and current_value > self.best_value
    ):
        self.best_value = current_value
        self.counter = 0
    else:
        self.counter += 1

    if self.counter >= self.patience:
        logger.info(
            f"EarlyStopping: No improvement in '{self.monitor}' for {self.patience} epochs. "
            "Stopping training."
        )
        trainer._stop_training.fill_(1)

`GradientMonitoring(on, called_every=1)`

Bases: Callback

Logs gradient statistics (e.g., mean, standard deviation, max) during training.

This callback monitors and logs statistics about the gradients of the model parameters to help debug or optimize the training process.

Example Usage in TrainConfig: To use GradientMonitoring, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import GradientMonitoring

# Create an instance of the GradientMonitoring callback
gradient_monitoring = GradientMonitoring(on="train_batch_end", called_every=10)

config = TrainConfig(
    max_iter=10000,
    print_every=1000,
    callbacks=[gradient_monitoring]
)

Initializes the GradientMonitoring callback.

PARAMETER	DESCRIPTION
`on`	The event to trigger the callback (e.g., "train_batch_end"). TYPE: `str`
`called_every`	Frequency of callback calls in terms of iterations. TYPE: `int` DEFAULT: `1`

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(self, on: str, called_every: int = 1):
    """Initializes the GradientMonitoring callback.

    Args:
        on (str): The event to trigger the callback (e.g., "train_batch_end").
        called_every (int): Frequency of callback calls in terms of iterations.
    """
    super().__init__(on=on, called_every=called_every)

`run_callback(trainer, config, writer)`

Logs gradient statistics.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
    """
    Logs gradient statistics.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter): The writer object for logging.
    """
    if trainer.accelerator.rank == 0:
        gradient_stats = {}
        for name, param in trainer.model.named_parameters():
            if param.grad is not None:
                grad = param.grad
                gradient_stats.update(
                    {
                        name + "_mean": grad.mean().item(),
                        name + "_std": grad.std().item(),
                        name + "_max": grad.max().item(),
                        name + "_min": grad.min().item(),
                    }
                )

        writer.write(trainer.opt_result.iteration, gradient_stats)

`LRSchedulerCosineAnnealing(on, called_every, t_max, min_lr=0.0)`

Bases: Callback

Applies cosine annealing to the learning rate during training.

This callback decreases the learning rate following a cosine curve, starting from the initial learning rate and annealing to a minimum (min_lr).

Example Usage in TrainConfig: To use LRSchedulerCosineAnnealing, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import LRSchedulerCosineAnnealing

# Create an instance of the LRSchedulerCosineAnnealing callback
lr_cosine = LRSchedulerCosineAnnealing(on="train_batch_end",
                                       called_every=1,
                                       t_max=5000,
                                       min_lr=1e-6)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback
    callbacks=[lr_cosine]
)

Initializes the LRSchedulerCosineAnnealing callback.

PARAMETER	DESCRIPTION
`on`	The event to trigger the callback. TYPE: `str`
`called_every`	Frequency of callback calls in terms of iterations. TYPE: `int`
`t_max`	The total number of iterations for one annealing cycle. TYPE: `int`
`min_lr`	The minimum learning rate. Default is 0.0. TYPE: `float` DEFAULT: `0.0`

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(self, on: str, called_every: int, t_max: int, min_lr: float = 0.0):
    """Initializes the LRSchedulerCosineAnnealing callback.

    Args:
        on (str): The event to trigger the callback.
        called_every (int): Frequency of callback calls in terms of iterations.
        t_max (int): The total number of iterations for one annealing cycle.
        min_lr (float, optional): The minimum learning rate. Default is 0.0.
    """
    super().__init__(on=on, called_every=called_every)
    self.t_max = t_max
    self.min_lr = min_lr

`run_callback(trainer, config, writer)`

Adjusts the learning rate using cosine annealing.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
    """
    Adjusts the learning rate using cosine annealing.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter): The writer object for logging.
    """
    for param_group in trainer.optimizer.param_groups:
        max_lr = param_group["lr"]
        new_lr = (
            self.min_lr
            + (max_lr - self.min_lr)
            * (1 + math.cos(math.pi * trainer.opt_result.iteration / self.t_max))
            / 2
        )
        param_group["lr"] = new_lr

`LRSchedulerCyclic(on, called_every, base_lr, max_lr, step_size)`

Bases: Callback

Applies a cyclic learning rate schedule during training.

This callback oscillates the learning rate between a minimum (base_lr) and a maximum (max_lr) over a defined cycle length (step_size). The learning rate follows a triangular wave pattern.

Example Usage in TrainConfig: To use LRSchedulerCyclic, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import LRSchedulerCyclic

# Create an instance of the LRSchedulerCyclic callback
lr_cyclic = LRSchedulerCyclic(on="train_batch_end",
                              called_every=1,
                              base_lr=0.001,
                              max_lr=0.01,
                              step_size=2000)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback
    callbacks=[lr_cyclic]
)

Initializes the LRSchedulerCyclic callback.

PARAMETER	DESCRIPTION
`on`	The event to trigger the callback. TYPE: `str`
`called_every`	Frequency of callback calls in terms of iterations. TYPE: `int`
`base_lr`	The minimum learning rate. TYPE: `float`
`max_lr`	The maximum learning rate. TYPE: `float`
`step_size`	Number of iterations for half a cycle. TYPE: `int`

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(self, on: str, called_every: int, base_lr: float, max_lr: float, step_size: int):
    """Initializes the LRSchedulerCyclic callback.

    Args:
        on (str): The event to trigger the callback.
        called_every (int): Frequency of callback calls in terms of iterations.
        base_lr (float): The minimum learning rate.
        max_lr (float): The maximum learning rate.
        step_size (int): Number of iterations for half a cycle.
    """
    super().__init__(on=on, called_every=called_every)
    self.base_lr = base_lr
    self.max_lr = max_lr
    self.step_size = step_size

`run_callback(trainer, config, writer)`

Adjusts the learning rate cyclically.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
    """
    Adjusts the learning rate cyclically.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter): The writer object for logging.
    """
    cycle = trainer.opt_result.iteration // (2 * self.step_size)
    x = abs(trainer.opt_result.iteration / self.step_size - 2 * cycle - 1)
    scale = max(0, (1 - x))
    new_lr = self.base_lr + (self.max_lr - self.base_lr) * scale
    for param_group in trainer.optimizer.param_groups:
        param_group["lr"] = new_lr

`LRSchedulerStepDecay(on, called_every, gamma=0.5)`

Bases: Callback

Reduces the learning rate by a factor at regular intervals.

This callback adjusts the learning rate by multiplying it with a decay factor after a specified number of iterations. The learning rate is updated as: lr = lr * gamma

Example Usage in TrainConfig: To use LRSchedulerStepDecay, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import LRSchedulerStepDecay

# Create an instance of the LRSchedulerStepDecay callback
lr_step_decay = LRSchedulerStepDecay(on="train_epoch_end",
                                     called_every=100,
                                     gamma=0.5)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback
    callbacks=[lr_step_decay]
)

Initializes the LRSchedulerStepDecay callback.

PARAMETER	DESCRIPTION
`on`	The event to trigger the callback. TYPE: `str`
`called_every`	Frequency of callback calls in terms of iterations. TYPE: `int`
`gamma`	The decay factor applied to the learning rate. A value < 1 reduces the learning rate over time. Default is 0.5. TYPE: `float` DEFAULT: `0.5`

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(self, on: str, called_every: int, gamma: float = 0.5):
    """Initializes the LRSchedulerStepDecay callback.

    Args:
        on (str): The event to trigger the callback.
        called_every (int): Frequency of callback calls in terms of iterations.
        gamma (float, optional): The decay factor applied to the learning rate.
            A value < 1 reduces the learning rate over time. Default is 0.5.
    """
    super().__init__(on=on, called_every=called_every)
    self.gamma = gamma

`run_callback(trainer, config, writer)`

Runs the callback to apply step decay to the learning rate.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
    """
    Runs the callback to apply step decay to the learning rate.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter): The writer object for logging.
    """
    for param_group in trainer.optimizer.param_groups:
        param_group["lr"] *= self.gamma

`LoadCheckpoint(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Bases: Callback

Callback to load a model checkpoint.

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`run_callback(trainer, config, writer)`

Loads a model checkpoint.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

RETURNS	DESCRIPTION
`Any`	The result of loading the checkpoint. TYPE: `Any`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Loads a model checkpoint.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.

    Returns:
        Any: The result of loading the checkpoint.
    """
    if trainer.accelerator.rank == 0:
        folder = config.log_folder
        model = trainer.model
        optimizer = trainer.optimizer
        device = trainer.accelerator.execution.log_device
        return load_checkpoint(folder, model, optimizer, device=device)

`LogHyperparameters(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Bases: Callback

Callback to log hyperparameters using the writer.

The LogHyperparameters callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use LogHyperparameters, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import LogHyperparameters

# Create an instance of the LogHyperparameters callback
log_hyper_callback = LogHyperparameters(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[log_hyper_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`run_callback(trainer, config, writer)`

Logs hyperparameters using the writer.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Logs hyperparameters using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    if trainer.accelerator.rank == 0:
        hyperparams = config.hyperparams
        writer.log_hyperparams(hyperparams)

`LogModelTracker(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Bases: Callback

Callback to log the model using the writer.

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`run_callback(trainer, config, writer)`

Logs the model using the writer.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Logs the model using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    if trainer.accelerator.rank == 0:
        model = trainer.model
        writer.log_model(
            model, trainer.train_dataloader, trainer.val_dataloader, trainer.test_dataloader
        )

`PlotMetrics(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Bases: Callback

Callback to plot metrics using the writer.

The PlotMetrics callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use PlotMetrics, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import PlotMetrics

# Create an instance of the PlotMetrics callback
plot_metrics_callback = PlotMetrics(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[plot_metrics_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`run_callback(trainer, config, writer)`

Plots metrics using the writer.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Plots metrics using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    if trainer.accelerator.rank == 0:
        opt_result = trainer.opt_result
        plotting_functions = config.plotting_functions
        writer.plot(trainer.model, opt_result.iteration, plotting_functions)

`PrintMetrics(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Bases: Callback

Callback to print metrics using the writer.

The PrintMetrics callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use PrintMetrics, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import PrintMetrics

# Create an instance of the PrintMetrics callback
print_metrics_callback = PrintMetrics(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[print_metrics_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`run_callback(trainer, config, writer)`

Prints metrics using the writer.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Prints metrics using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    opt_result = trainer.opt_result
    writer.print_metrics(opt_result)

`SaveBestCheckpoint(on, called_every)`

Bases: SaveCheckpoint

Callback to save the best model checkpoint based on a validation criterion.

Initializes the SaveBestCheckpoint callback.

PARAMETER	DESCRIPTION
`on`	The event to trigger the callback. TYPE: `str`
`called_every`	Frequency of callback calls in terms of iterations. TYPE: `int`

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(self, on: str, called_every: int):
    """Initializes the SaveBestCheckpoint callback.

    Args:
        on (str): The event to trigger the callback.
        called_every (int): Frequency of callback calls in terms of iterations.
    """
    super().__init__(on=on, called_every=called_every)
    self.best_loss = float("inf")

`run_callback(trainer, config, writer)`

Saves the checkpoint if the current loss is better than the best loss.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Saves the checkpoint if the current loss is better than the best loss.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    if trainer.accelerator.rank == 0:
        opt_result = trainer.opt_result
        if config.validation_criterion and config.validation_criterion(
            opt_result.loss, self.best_loss, config.val_epsilon
        ):
            self.best_loss = opt_result.loss

            folder = config.log_folder
            model = trainer.model
            optimizer = trainer.optimizer
            opt_result = trainer.opt_result
            write_checkpoint(folder, model, optimizer, "best")

`SaveCheckpoint(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Bases: Callback

Callback to save a model checkpoint.

The SaveCheckpoint callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use SaveCheckpoint, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import SaveCheckpoint

# Create an instance of the SaveCheckpoint callback
save_checkpoint_callback = SaveCheckpoint(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[save_checkpoint_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`run_callback(trainer, config, writer)`

Saves a model checkpoint.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Saves a model checkpoint.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    if trainer.accelerator.rank == 0:
        folder = config.log_folder
        model = trainer.model
        optimizer = trainer.optimizer
        opt_result = trainer.opt_result
        write_checkpoint(folder, model, optimizer, opt_result.iteration)

`WriteMetrics(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)`

Bases: Callback

Callback to write metrics using the writer.

The WriteMetrics callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use WriteMetrics, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import WriteMetrics

# Create an instance of the WriteMetrics callback
write_metrics_callback = WriteMetrics(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[write_metrics_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = (
        callback_condition if callback_condition else Callback.default_callback
    )

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
            opt_res, modify_optimize_result
        )
    else:
        self.modify_optimize_result = (
            modify_optimize_result
            if modify_optimize_result
            else Callback.modify_opt_res_default
        )

`run_callback(trainer, config, writer)`

Writes metrics using the writer.

PARAMETER	DESCRIPTION
`trainer`	The training object. TYPE: `Any`
`config`	The configuration object. TYPE: `TrainConfig`
`writer`	The writer object for logging. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Writes metrics using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    if trainer.accelerator.rank == 0:
        opt_result = trainer.opt_result
        writer.write(opt_result.iteration, opt_result.metrics)

`BaseTrainer(model, optimizer, config, loss_fn='mse', optimize_step=optimize_step, train_dataloader=None, val_dataloader=None, test_dataloader=None, max_batches=None)`

Base class for training machine learning models using a given optimizer.

The base class implements contextmanager for gradient based/free optimization, properties, property setters, input validations, callback decorator generator, and empty hooks for different training steps.

This class provides

Context managers for enabling/disabling gradient-based optimization
Properties for managing models, optimizers, and dataloaders
Input validations and a callback decorator generator
Config and callback managers using the provided TrainConfig

ATTRIBUTE	DESCRIPTION
`use_grad`	Indicates if gradients are used for optimization. Default is True. TYPE: `bool`
`model`	The neural network model. TYPE: `Module`
`optimizer`	The optimizer for training. TYPE: `Optimizer \| Optimizer \| None`
`config`	The configuration settings for training. TYPE: `TrainConfig`
`train_dataloader`	DataLoader for training data. TYPE: `Dataloader \| DictDataLoader \| None`
`val_dataloader`	DataLoader for validation data. TYPE: `Dataloader \| DictDataLoader \| None`
`test_dataloader`	DataLoader for testing data. TYPE: `Dataloader \| DictDataLoader \| None`
`optimize_step`	Function for performing an optimization step. TYPE: `Callable`
`loss_fn`	loss function to use. Default loss function used is 'mse' TYPE: `Callable \| str ]`
`num_training_batches`	Number of training batches. In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: `int`
`num_validation_batches`	Number of validation batches. In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: `int`
`num_test_batches`	Number of test batches. In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: `int`
`state`	Current state in the training process TYPE: `str`

Initializes the BaseTrainer.

PARAMETER	DESCRIPTION
`model`	The model to train. TYPE: `Module`
`optimizer`	The optimizer for training. TYPE: `Optimizer \| Optimizer \| None`
`config`	The TrainConfig settings for training. TYPE: `TrainConfig`
`loss_fn`	The loss function to use. str input to be specified to use a default loss function. currently supported loss functions: 'mse', 'cross_entropy'. If not specified, default mse loss will be used. TYPE: `str \| Callable` DEFAULT: `'mse'`
`train_dataloader`	DataLoader for training data. If the model does not need data to evaluate loss, no dataset should be provided. TYPE: `Dataloader \| DictDataLoader \| None` DEFAULT: `None`
`val_dataloader`	DataLoader for validation data. TYPE: `Dataloader \| DictDataLoader \| None` DEFAULT: `None`
`test_dataloader`	DataLoader for testing data. TYPE: `Dataloader \| DictDataLoader \| None` DEFAULT: `None`
`max_batches`	Maximum number of batches to process per epoch. This is only valid in case of finite TensorDataset dataloaders. if max_batches is not None, the maximum number of batches used will be min(max_batches, len(dataloader.dataset)) In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: `int \| None` DEFAULT: `None`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def __init__(
    self,
    model: nn.Module,
    optimizer: optim.Optimizer | NGOptimizer | None,
    config: TrainConfig,
    loss_fn: str | Callable = "mse",
    optimize_step: Callable = optimize_step,
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
    test_dataloader: DataLoader | DictDataLoader | None = None,
    max_batches: int | None = None,
):
    """
    Initializes the BaseTrainer.

    Args:
        model (nn.Module): The model to train.
        optimizer (optim.Optimizer | NGOptimizer | None): The optimizer
            for training.
        config (TrainConfig): The TrainConfig settings for training.
        loss_fn (str | Callable): The loss function to use.
            str input to be specified to use a default loss function.
            currently supported loss functions: 'mse', 'cross_entropy'.
            If not specified, default mse loss will be used.
        train_dataloader (Dataloader | DictDataLoader | None): DataLoader for training data.
            If the model does not need data to evaluate loss, no dataset
            should be provided.
        val_dataloader (Dataloader | DictDataLoader | None): DataLoader for validation data.
        test_dataloader (Dataloader | DictDataLoader | None): DataLoader for testing data.
        max_batches (int | None): Maximum number of batches to process per epoch.
            This is only valid in case of finite TensorDataset dataloaders.
            if max_batches is not None, the maximum number of batches used will
            be min(max_batches, len(dataloader.dataset))
            In case of InfiniteTensorDataset only 1 batch per epoch is used.
    """
    self._model: nn.Module
    self._optimizer: optim.Optimizer | NGOptimizer | None
    self._config: TrainConfig
    self._train_dataloader: DataLoader | DictDataLoader | None = None
    self._val_dataloader: DataLoader | DictDataLoader | None = None
    self._test_dataloader: DataLoader | DictDataLoader | None = None

    self.config = config
    self.model = model
    self.optimizer = optimizer
    self.max_batches = max_batches

    self.num_training_batches: int
    self.num_validation_batches: int
    self.num_test_batches: int

    self.train_dataloader = train_dataloader
    self.val_dataloader = val_dataloader
    self.test_dataloader = test_dataloader

    self.loss_fn: Callable = get_loss_fn(loss_fn)
    self.optimize_step: Callable = optimize_step
    self.ng_params: ng.p.Array
    self.training_stage: TrainingStage = TrainingStage("idle")

`config` `property` `writable`

Returns the training configuration.

RETURNS	DESCRIPTION
`TrainConfig`	The configuration object. TYPE: `TrainConfig`

`model` `property` `writable`

Returns the model if set, otherwise raises an error.

RETURNS	DESCRIPTION
`Module`	nn.Module: The model.

`optimizer` `property` `writable`

Returns the optimizer if set, otherwise raises an error.

RETURNS	DESCRIPTION
`Optimizer \| Optimizer \| None`	optim.Optimizer \| NGOptimizer \| None: The optimizer.

`test_dataloader` `property` `writable`

Returns the test DataLoader, validating its type.

RETURNS	DESCRIPTION
`DataLoader`	The DataLoader for testing data. TYPE: `DataLoader`

`train_dataloader` `property` `writable`

Returns the training DataLoader, validating its type.

RETURNS	DESCRIPTION
`DataLoader`	The DataLoader for training data. TYPE: `DataLoader`

`use_grad` `property` `writable`

Returns the optimization framework for the trainer.

use_grad = True : Gradient based optimization use_grad = False : Gradient free optimization

RETURNS	DESCRIPTION
`bool`	Bool value for using gradient. TYPE: `bool`

`val_dataloader` `property` `writable`

Returns the validation DataLoader, validating its type.

RETURNS	DESCRIPTION
`DataLoader`	The DataLoader for validation data. TYPE: `DataLoader`

`_compute_num_batches(dataloader)`

Computes the number of batches for the given DataLoader.

PARAMETER	DESCRIPTION
`dataloader`	The DataLoader for which to compute the number of batches. TYPE: `DataLoader`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def _compute_num_batches(self, dataloader: DataLoader | DictDataLoader) -> int:
    """
    Computes the number of batches for the given DataLoader.

    Args:
        dataloader (DataLoader): The DataLoader for which to compute
            the number of batches.
    """
    if dataloader is None:
        return 1
    if isinstance(dataloader, DictDataLoader):
        dataloader_name, dataloader_value = list(dataloader.dataloaders.items())[0]
        dataset = dataloader_value.dataset
        batch_size = dataloader_value.batch_size
    else:
        dataset = dataloader.dataset
        batch_size = dataloader.batch_size

    if isinstance(dataset, TensorDataset):
        n_batches = int((dataset.tensors[0].size(0) + batch_size - 1) // batch_size)
        return min(self.max_batches, n_batches) if self.max_batches is not None else n_batches
    else:
        return 1

`_validate_dataloader(dataloader, dataloader_type)`

Validates the type of the DataLoader and raises errors for unsupported types.

PARAMETER	DESCRIPTION
`dataloader`	The DataLoader to validate. TYPE: `DataLoader \| DictDataLoader`
`dataloader_type`	The type of DataLoader ("train", "val", or "test"). TYPE: `str`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def _validate_dataloader(
    self, dataloader: DataLoader | DictDataLoader, dataloader_type: str
) -> None:
    """
    Validates the type of the DataLoader and raises errors for unsupported types.

    Args:
        dataloader (DataLoader | DictDataLoader): The DataLoader to validate.
        dataloader_type (str): The type of DataLoader ("train", "val", or "test").
    """
    if dataloader is not None:
        if not isinstance(dataloader, (DataLoader, DictDataLoader)):
            raise NotImplementedError(
                f"Unsupported dataloader type: {type(dataloader)}."
                "The dataloader must be an instance of DataLoader."
            )
    if dataloader_type == "val" and self.config.val_every > 0:
        if not isinstance(dataloader, (DataLoader, DictDataLoader)):
            raise ValueError(
                "If `config.val_every` is provided as an integer > 0, validation_dataloader"
                "must be an instance of `DataLoader` or `DictDataLoader`."
            )

`callback(phase)` `staticmethod`

Decorator for executing callbacks before and after a phase.

Phase are different hooks during the training. list of valid phases is defined in Callbacks. We also update the current state of the training process in the callback decorator.

PARAMETER	DESCRIPTION
`phase`	The phase for which the callback is executed (e.g., "train", "train_epoch", "train_batch"). TYPE: `str`

RETURNS	DESCRIPTION
`Callable`	The decorated function. TYPE: `Callable`

Source code in qadence/ml_tools/train_utils/base_trainer.py

@staticmethod
def callback(phase: str) -> Callable:
    """
    Decorator for executing callbacks before and after a phase.

    Phase are different hooks during the training. list of valid
    phases is defined in Callbacks.
    We also update the current state of the training process in
    the callback decorator.

    Args:
        phase (str): The phase for which the callback is executed (e.g., "train",
            "train_epoch", "train_batch").

    Returns:
        Callable: The decorated function.
    """

    def decorator(method: Callable) -> Callable:
        def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
            start_event = f"{phase}_start"
            end_event = f"{phase}_end"

            self.training_stage = TrainingStage(start_event)
            self.callback_manager.run_callbacks(trainer=self)
            result = method(self, *args, **kwargs)

            self.training_stage = TrainingStage(end_event)
            # build_optimize_result method is defined in the trainer.
            self.build_optimize_result(result)
            self.callback_manager.run_callbacks(trainer=self)

            return result

        return wrapper

    return decorator

`disable_grad_opt(optimizer=None)`

Context manager to temporarily disable gradient-based optimization.

PARAMETER	DESCRIPTION
`optimizer`	The Nevergrad optimizer to use. If no optimizer is provided, default optimizer for trainer object will be used. TYPE: `Optimizer` DEFAULT: `None`

Source code in qadence/ml_tools/train_utils/base_trainer.py

@contextmanager
def disable_grad_opt(self, optimizer: NGOptimizer | None = None) -> Iterator[None]:
    """
    Context manager to temporarily disable gradient-based optimization.

    Args:
        optimizer (NGOptimizer): The Nevergrad optimizer to use.
            If no optimizer is provided, default optimizer for trainer
            object will be used.
    """
    original_mode = self.use_grad
    original_optimizer = self._optimizer
    try:
        self.use_grad = False
        self.callback_manager.use_grad = False
        self.optimizer = optimizer if optimizer else self.optimizer
        yield
    finally:
        self.use_grad = original_mode
        self.callback_manager.use_grad = original_mode
        self.optimizer = original_optimizer

`enable_grad_opt(optimizer=None)`

Context manager to temporarily enable gradient-based optimization.

PARAMETER	DESCRIPTION
`optimizer`	The PyTorch optimizer to use. If no optimizer is provided, default optimizer for trainer object will be used. TYPE: `Optimizer` DEFAULT: `None`

Source code in qadence/ml_tools/train_utils/base_trainer.py

@contextmanager
def enable_grad_opt(self, optimizer: optim.Optimizer | None = None) -> Iterator[None]:
    """
    Context manager to temporarily enable gradient-based optimization.

    Args:
        optimizer (optim.Optimizer): The PyTorch optimizer to use.
            If no optimizer is provided, default optimizer for trainer
            object will be used.
    """
    original_mode = self.use_grad
    original_optimizer = self._optimizer
    try:
        self.use_grad = True
        self.callback_manager.use_grad = True
        self.optimizer = optimizer if optimizer else self.optimizer
        yield
    finally:
        self.use_grad = original_mode
        self.callback_manager.use_grad = original_mode
        self.optimizer = original_optimizer

`on_test_batch_end(test_batch_loss_metrics)`

Called at the end of each testing batch.

PARAMETER	DESCRIPTION
`test_batch_loss_metrics`	Metrics for the testing batch loss. tuple of (loss, metrics) TYPE: `tuple[Tensor, Any]`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_test_batch_end(self, test_batch_loss_metrics: tuple[torch.Tensor, Any]) -> None:
    """
    Called at the end of each testing batch.

    Args:
        test_batch_loss_metrics: Metrics for the testing batch loss.
            tuple of (loss, metrics)
    """
    pass

`on_test_batch_start(batch)`

Called at the start of each testing batch.

PARAMETER	DESCRIPTION
`batch`	A batch of data from the DataLoader. Typically a tuple containing input tensors and corresponding target tensors. TYPE: `tuple[Tensor, ...] \| None`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_test_batch_start(self, batch: tuple[torch.Tensor, ...] | None) -> None:
    """
    Called at the start of each testing batch.

    Args:
        batch: A batch of data from the DataLoader. Typically a tuple containing
            input tensors and corresponding target tensors.
    """
    pass

`on_train_batch_end(train_batch_loss_metrics)`

Called at the end of each training batch.

PARAMETER	DESCRIPTION
`train_batch_loss_metrics`	Metrics for the training batch loss. tuple of (loss, metrics) TYPE: `tuple[Tensor, Any]`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_batch_end(self, train_batch_loss_metrics: tuple[torch.Tensor, Any]) -> None:
    """
    Called at the end of each training batch.

    Args:
        train_batch_loss_metrics: Metrics for the training batch loss.
            tuple of (loss, metrics)
    """
    pass

`on_train_batch_start(batch)`

Called at the start of each training batch.

PARAMETER	DESCRIPTION
`batch`	A batch of data from the DataLoader. Typically a tuple containing input tensors and corresponding target tensors. TYPE: `tuple[Tensor, ...] \| None`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_batch_start(self, batch: tuple[torch.Tensor, ...] | None) -> None:
    """
    Called at the start of each training batch.

    Args:
        batch: A batch of data from the DataLoader. Typically a tuple containing
            input tensors and corresponding target tensors.
    """
    pass

`on_train_end(train_losses, val_losses=None)`

Called at the end of training.

PARAMETER	DESCRIPTION
`train_losses`	Metrics for the training losses. list -> list -> tuples Epochs -> Training Batches -> (loss, metrics) TYPE: `list[list[tuple[Tensor, Any]]]`
`val_losses`	Metrics for the validation losses. list -> list -> tuples Epochs -> Validation Batches -> (loss, metrics) TYPE: `list[list[tuple[Tensor, Any]]] \| None` DEFAULT: `None`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_end(
    self,
    train_losses: list[list[tuple[torch.Tensor, Any]]],
    val_losses: list[list[tuple[torch.Tensor, Any]]] | None = None,
) -> None:
    """
    Called at the end of training.

    Args:
        train_losses (list[list[tuple[torch.Tensor, Any]]]):
            Metrics for the training losses.
            list    -> list                  -> tuples
            Epochs  -> Training Batches      -> (loss, metrics)
        val_losses (list[list[tuple[torch.Tensor, Any]]] | None):
            Metrics for the validation losses.
            list    -> list                  -> tuples
            Epochs  -> Validation Batches    -> (loss, metrics)
    """
    pass

`on_train_epoch_end(train_epoch_loss_metrics)`

Called at the end of each training epoch.

PARAMETER	DESCRIPTION
`train_epoch_loss_metrics`	Metrics for the training epoch losses. list -> tuples Training Batches -> (loss, metrics) TYPE: `list[tuple[Tensor, Any]]`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_epoch_end(self, train_epoch_loss_metrics: list[tuple[torch.Tensor, Any]]) -> None:
    """
    Called at the end of each training epoch.

    Args:
        train_epoch_loss_metrics: Metrics for the training epoch losses.
            list                  -> tuples
            Training Batches      -> (loss, metrics)
    """
    pass

`on_train_epoch_start()`

Called at the start of each training epoch.

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_epoch_start(self) -> None:
    """Called at the start of each training epoch."""
    pass

`on_train_start()`

Called at the start of training.

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_start(self) -> None:
    """Called at the start of training."""
    pass

`on_val_batch_end(val_batch_loss_metrics)`

Called at the end of each validation batch.

PARAMETER	DESCRIPTION
`val_batch_loss_metrics`	Metrics for the validation batch loss. tuple of (loss, metrics) TYPE: `tuple[Tensor, Any]`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_batch_end(self, val_batch_loss_metrics: tuple[torch.Tensor, Any]) -> None:
    """
    Called at the end of each validation batch.

    Args:
        val_batch_loss_metrics: Metrics for the validation batch loss.
            tuple of (loss, metrics)
    """
    pass

`on_val_batch_start(batch)`

Called at the start of each validation batch.

PARAMETER	DESCRIPTION
`batch`	A batch of data from the DataLoader. Typically a tuple containing input tensors and corresponding target tensors. TYPE: `tuple[Tensor, ...] \| None`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_batch_start(self, batch: tuple[torch.Tensor, ...] | None) -> None:
    """
    Called at the start of each validation batch.

    Args:
        batch: A batch of data from the DataLoader. Typically a tuple containing
            input tensors and corresponding target tensors.
    """
    pass

`on_val_epoch_end(val_epoch_loss_metrics)`

Called at the end of each validation epoch.

PARAMETER	DESCRIPTION
`val_epoch_loss_metrics`	Metrics for the validation epoch loss. list -> tuples Validation Batches -> (loss, metrics) TYPE: `list[tuple[Tensor, Any]]`

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_epoch_end(self, val_epoch_loss_metrics: list[tuple[torch.Tensor, Any]]) -> None:
    """
    Called at the end of each validation epoch.

    Args:
        val_epoch_loss_metrics: Metrics for the validation epoch loss.
            list                    -> tuples
            Validation Batches      -> (loss, metrics)
    """
    pass

`on_val_epoch_start()`

Called at the start of each validation epoch.

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_epoch_start(self) -> None:
    """Called at the start of each validation epoch."""
    pass

`set_use_grad(value)` `classmethod`

Sets the global use_grad flag.

PARAMETER	DESCRIPTION
`value`	Whether to use gradient-based optimization. TYPE: `bool`

Source code in qadence/ml_tools/train_utils/base_trainer.py

@classmethod
def set_use_grad(cls, value: bool) -> None:
    """
    Sets the global use_grad flag.

    Args:
        value (bool): Whether to use gradient-based optimization.
    """
    if not isinstance(value, bool):
        raise TypeError("use_grad must be a boolean value.")
    cls._use_grad = value

`BaseWriter`

Bases: ABC

Abstract base class for experiment tracking writers.

METHOD	DESCRIPTION
`open`	Opens the writer and sets up the logging environment.
`close`	Closes the writer and finalizes any ongoing logging processes.
`print_metrics`	Prints metrics and loss in a formatted manner.
`write`	Writes the optimization results to the tracking tool.
`log_hyperparams`	Logs the hyperparameters to the tracking tool.
`plot`	Logs model plots using provided plotting functions.
`log_model`	Logs the model and any relevant information.

`close()` `abstractmethod`

Closes the writer and finalizes logging.

Source code in qadence/ml_tools/callbacks/writer_registry.py

@abstractmethod
def close(self) -> None:
    """Closes the writer and finalizes logging."""
    raise NotImplementedError("Writers must implement a close method.")

`log_hyperparams(hyperparams)` `abstractmethod`

Logs hyperparameters.

PARAMETER	DESCRIPTION
`hyperparams`	A dictionary of hyperparameters to log. TYPE: `dict`

Source code in qadence/ml_tools/callbacks/writer_registry.py

@abstractmethod
def log_hyperparams(self, hyperparams: dict) -> None:
    """
    Logs hyperparameters.

    Args:
        hyperparams (dict): A dictionary of hyperparameters to log.
    """
    raise NotImplementedError("Writers must implement a log_hyperparams method.")

`log_model(model, train_dataloader=None, val_dataloader=None, test_dataloader=None)` `abstractmethod`

Logs the model and associated data.

PARAMETER	DESCRIPTION
`model`	The model to log. TYPE: `Module`
`train_dataloader`	DataLoader for training data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`val_dataloader`	DataLoader for validation data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`test_dataloader`	DataLoader for testing data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`

Source code in qadence/ml_tools/callbacks/writer_registry.py

@abstractmethod
def log_model(
    self,
    model: Module,
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
    test_dataloader: DataLoader | DictDataLoader | None = None,
) -> None:
    """
    Logs the model and associated data.

    Args:
        model (Module): The model to log.
        train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
        val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.
        test_dataloader (DataLoader | DictDataLoader |  None): DataLoader for testing data.
    """
    raise NotImplementedError("Writers must implement a log_model method.")

`open(config, iteration=None)` `abstractmethod`

Opens the writer and prepares it for logging.

PARAMETER	DESCRIPTION
`config`	Configuration object containing settings for logging. TYPE: `TrainConfig`
`iteration`	The iteration step to start logging from. Defaults to None. TYPE: `int` DEFAULT: `None`

Source code in qadence/ml_tools/callbacks/writer_registry.py

@abstractmethod
def open(self, config: TrainConfig, iteration: int | None = None) -> Any:
    """
    Opens the writer and prepares it for logging.

    Args:
        config: Configuration object containing settings for logging.
        iteration (int, optional): The iteration step to start logging from.
            Defaults to None.
    """
    raise NotImplementedError("Writers must implement an open method.")

`plot(model, iteration, plotting_functions)` `abstractmethod`

Logs plots of the model using provided plotting functions.

PARAMETER	DESCRIPTION
`model`	The model to plot. TYPE: `Module`
`iteration`	The current iteration number. TYPE: `int`
`plotting_functions`	Functions used to generate plots. TYPE: `tuple[PlottingFunction, ...]`

Source code in qadence/ml_tools/callbacks/writer_registry.py

@abstractmethod
def plot(
    self,
    model: Module,
    iteration: int,
    plotting_functions: tuple[PlottingFunction, ...],
) -> None:
    """
    Logs plots of the model using provided plotting functions.

    Args:
        model (Module): The model to plot.
        iteration (int): The current iteration number.
        plotting_functions (tuple[PlottingFunction, ...]): Functions used to
            generate plots.
    """
    raise NotImplementedError("Writers must implement a plot method.")

`print_metrics(result)`

Prints the metrics and loss in a readable format.

PARAMETER	DESCRIPTION
`result`	The optimization results to display. TYPE: `OptimizeResult`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def print_metrics(self, result: OptimizeResult) -> None:
    """Prints the metrics and loss in a readable format.

    Args:
        result (OptimizeResult): The optimization results to display.
    """

    # Find the key in result.metrics that contains "loss" (case-insensitive)
    loss_key = next((k for k in result.metrics if "loss" in k.lower()), None)
    initial = f"P {result.rank: >2}|{result.device: <7}| Iteration {result.iteration: >7}| "
    if loss_key:
        loss_value = result.metrics[loss_key]
        msg = initial + f"{loss_key.title()}: {loss_value:.7f} -"
    else:
        msg = initial + f"Loss: None -"
    msg += " ".join([f"{k}: {v:.7f}" for k, v in result.metrics.items() if k != loss_key])
    print(msg)

`write(iteration, metrics)` `abstractmethod`

Logs the results of the current iteration.

PARAMETER	DESCRIPTION
`iteration`	The current training iteration. TYPE: `int`
`metrics`	A dictionary of metrics to log, where keys are metric names and values are the corresponding metric values. TYPE: `dict`

Source code in qadence/ml_tools/callbacks/writer_registry.py

@abstractmethod
def write(self, iteration: int, metrics: dict) -> None:
    """
    Logs the results of the current iteration.

    Args:
        iteration (int): The current training iteration.
        metrics (dict): A dictionary of metrics to log, where keys are metric names
                        and values are the corresponding metric values.
    """
    raise NotImplementedError("Writers must implement a write method.")

`MLFlowWriter()`

Bases: BaseWriter

Writer for logging to MLflow.

ATTRIBUTE	DESCRIPTION
`run`	The active MLflow run. TYPE: `Run`
`mlflow`	The MLflow module. TYPE: `ModuleType`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def __init__(self) -> None:
    try:
        from mlflow.entities import Run
    except ImportError:
        raise ImportError(
            "mlflow is not installed. Please install qadence with the mlflow feature: "
            "`pip install qadence[mlflow]`."
        )

    self.run: Run
    self.mlflow: ModuleType

`close()`

Closes the MLflow run.

Source code in qadence/ml_tools/callbacks/writer_registry.py

def close(self) -> None:
    """Closes the MLflow run."""
    if self.run:
        self.mlflow.end_run()

`get_signature_from_dataloader(model, dataloader)`

Infers the signature of the model based on the input data from the dataloader.

PARAMETER	DESCRIPTION
`model`	The model to use for inference. TYPE: `Module`
`dataloader`	DataLoader for model inputs. TYPE: `DataLoader \| DictDataLoader \| None`

RETURNS	DESCRIPTION
`Any`	Optional[Any]: The inferred signature, if available.

Source code in qadence/ml_tools/callbacks/writer_registry.py

def get_signature_from_dataloader(
    self, model: Module, dataloader: DataLoader | DictDataLoader | None
) -> Any:
    """
    Infers the signature of the model based on the input data from the dataloader.

    Args:
        model (Module): The model to use for inference.
        dataloader (DataLoader | DictDataLoader |  None): DataLoader for model inputs.

    Returns:
        Optional[Any]: The inferred signature, if available.
    """
    from mlflow.models import infer_signature

    if dataloader is None:
        return None

    xs: InputData
    xs, *_ = next(iter(dataloader))
    preds = model(xs)

    if isinstance(xs, Tensor):
        xs = xs.detach().cpu().numpy()
        preds = preds.detach().cpu().numpy()
        return infer_signature(xs, preds)

    return None

`log_hyperparams(hyperparams)`

Logs hyperparameters to MLflow.

PARAMETER	DESCRIPTION
`hyperparams`	A dictionary of hyperparameters to log. TYPE: `dict`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def log_hyperparams(self, hyperparams: dict) -> None:
    """
    Logs hyperparameters to MLflow.

    Args:
        hyperparams (dict): A dictionary of hyperparameters to log.
    """
    if self.mlflow:
        self.mlflow.log_params(hyperparams)
    else:
        raise RuntimeError(
            "The writer is not initialized."
            "Please call the 'writer.open()' method before writing"
        )

`log_model(model, train_dataloader=None, val_dataloader=None, test_dataloader=None)`

Logs the model and its signature to MLflow using the provided data loaders.

PARAMETER	DESCRIPTION
`model`	The model to log. TYPE: `Module`
`train_dataloader`	DataLoader for training data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`val_dataloader`	DataLoader for validation data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`test_dataloader`	DataLoader for testing data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def log_model(
    self,
    model: Module,
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
    test_dataloader: DataLoader | DictDataLoader | None = None,
) -> None:
    """
    Logs the model and its signature to MLflow using the provided data loaders.

    Args:
        model (Module): The model to log.
        train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
        val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.
        test_dataloader (DataLoader | DictDataLoader |  None): DataLoader for testing data.
    """
    if not self.mlflow:
        raise RuntimeError(
            "The writer is not initialized."
            "Please call the 'writer.open()' method before writing"
        )

    signatures = self.get_signature_from_dataloader(model, train_dataloader)
    self.mlflow.pytorch.log_model(model, artifact_path="model", signature=signatures)

`open(config, iteration=None)`

Opens the MLflow writer and initializes an MLflow run.

PARAMETER	DESCRIPTION
`config`	Configuration object containing settings for logging. TYPE: `TrainConfig`
`iteration`	The iteration step to start logging from. Defaults to None. TYPE: `int` DEFAULT: `None`

RETURNS	DESCRIPTION
`mlflow`	The MLflow module instance. TYPE: `ModuleType \| None`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def open(self, config: TrainConfig, iteration: int | None = None) -> ModuleType | None:
    """
    Opens the MLflow writer and initializes an MLflow run.

    Args:
        config: Configuration object containing settings for logging.
        iteration (int, optional): The iteration step to start logging from.
            Defaults to None.

    Returns:
        mlflow: The MLflow module instance.
    """
    import mlflow

    self.mlflow = mlflow
    tracking_uri = os.getenv("MLFLOW_TRACKING_URI", "")
    experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", str(uuid4()))
    run_name = os.getenv("MLFLOW_RUN_NAME", str(uuid4()))

    if self.mlflow:
        self.mlflow.set_tracking_uri(tracking_uri)

        # Create or get the experiment
        exp_filter_string = f"name = '{experiment_name}'"
        experiments = self.mlflow.search_experiments(filter_string=exp_filter_string)
        if not experiments:
            self.mlflow.create_experiment(name=experiment_name)

        self.mlflow.set_experiment(experiment_name)
        self.run = self.mlflow.start_run(run_name=run_name, nested=False)

    return self.mlflow

`plot(model, iteration, plotting_functions)`

Logs plots of the model using provided plotting functions.

PARAMETER	DESCRIPTION
`model`	The model to plot. TYPE: `Module`
`iteration`	The current iteration number. TYPE: `int`
`plotting_functions`	Functions used to generate plots. TYPE: `tuple[PlottingFunction, ...]`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def plot(
    self,
    model: Module,
    iteration: int,
    plotting_functions: tuple[PlottingFunction, ...],
) -> None:
    """
    Logs plots of the model using provided plotting functions.

    Args:
        model (Module): The model to plot.
        iteration (int): The current iteration number.
        plotting_functions (tuple[PlottingFunction, ...]): Functions used
            to generate plots.
    """
    if self.mlflow:
        for pf in plotting_functions:
            descr, fig = pf(model, iteration)
            self.mlflow.log_figure(fig, descr)
    else:
        raise RuntimeError(
            "The writer is not initialized."
            "Please call the 'writer.open()' method before writing"
        )

`write(iteration, metrics)`

Logs the results of the current iteration to MLflow.

PARAMETER	DESCRIPTION
`iteration`	The current training iteration. TYPE: `int`
`metrics`	A dictionary of metrics to log, where keys are metric names and values are the corresponding metric values. TYPE: `dict`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def write(self, iteration: int, metrics: dict) -> None:
    """
    Logs the results of the current iteration to MLflow.

    Args:
        iteration (int): The current training iteration.
        metrics (dict): A dictionary of metrics to log, where keys are metric names
                        and values are the corresponding metric values.
    """
    if self.mlflow:
        self.mlflow.log_metrics(metrics, step=iteration)
    else:
        raise RuntimeError(
            "The writer is not initialized."
            "Please call the 'writer.open()' method before writing."
        )

`TensorBoardWriter()`

Bases: BaseWriter

Writer for logging to TensorBoard.

ATTRIBUTE	DESCRIPTION
`writer`	The TensorBoard SummaryWriter instance. TYPE: `SummaryWriter`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def __init__(self) -> None:
    self.writer = None

`close()`

Closes the TensorBoard writer.

Source code in qadence/ml_tools/callbacks/writer_registry.py

def close(self) -> None:
    """Closes the TensorBoard writer."""
    if self.writer:
        self.writer.close()

`log_hyperparams(hyperparams)`

Logs hyperparameters to TensorBoard.

PARAMETER	DESCRIPTION
`hyperparams`	A dictionary of hyperparameters to log. TYPE: `dict`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def log_hyperparams(self, hyperparams: dict) -> None:
    """
    Logs hyperparameters to TensorBoard.

    Args:
        hyperparams (dict): A dictionary of hyperparameters to log.
    """
    if self.writer:
        self.writer.add_hparams(hyperparams, {})
    else:
        raise RuntimeError(
            "The writer is not initialized."
            "Please call the 'writer.open()' method before writing"
        )

`log_model(model, train_dataloader=None, val_dataloader=None, test_dataloader=None)`

Logs the model.

Currently not supported by TensorBoard.

PARAMETER	DESCRIPTION
`model`	The model to log. TYPE: `Module`
`train_dataloader`	DataLoader for training data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`val_dataloader`	DataLoader for validation data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`
`test_dataloader`	DataLoader for testing data. TYPE: `DataLoader \| DictDataLoader \| None` DEFAULT: `None`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def log_model(
    self,
    model: Module,
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
    test_dataloader: DataLoader | DictDataLoader | None = None,
) -> None:
    """
    Logs the model.

    Currently not supported by TensorBoard.

    Args:
        model (Module): The model to log.
        train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
        val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.
        test_dataloader (DataLoader | DictDataLoader |  None): DataLoader for testing data.
    """
    logger.warning("Model logging is not supported by tensorboard. No model will be logged.")

`open(config, iteration=None)`

Opens the TensorBoard writer.

PARAMETER	DESCRIPTION
`config`	Configuration object containing settings for logging. TYPE: `TrainConfig`
`iteration`	The iteration step to start logging from. Defaults to None. TYPE: `int` DEFAULT: `None`

RETURNS	DESCRIPTION
`SummaryWriter`	The initialized TensorBoard writer. TYPE: `SummaryWriter`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def open(self, config: TrainConfig, iteration: int | None = None) -> SummaryWriter:
    """
    Opens the TensorBoard writer.

    Args:
        config: Configuration object containing settings for logging.
        iteration (int, optional): The iteration step to start logging from.
            Defaults to None.

    Returns:
        SummaryWriter: The initialized TensorBoard writer.
    """
    log_dir = str(config.log_folder)
    purge_step = iteration if isinstance(iteration, int) else None
    self.writer = SummaryWriter(log_dir=log_dir, purge_step=purge_step)
    return self.writer

`plot(model, iteration, plotting_functions)`

Logs plots of the model using provided plotting functions.

PARAMETER	DESCRIPTION
`model`	The model to plot. TYPE: `Module`
`iteration`	The current iteration number. TYPE: `int`
`plotting_functions`	Functions used to generate plots. TYPE: `tuple[PlottingFunction, ...]`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def plot(
    self,
    model: Module,
    iteration: int,
    plotting_functions: tuple[PlottingFunction, ...],
) -> None:
    """
    Logs plots of the model using provided plotting functions.

    Args:
        model (Module): The model to plot.
        iteration (int): The current iteration number.
        plotting_functions (tuple[PlottingFunction, ...]): Functions used
            to generate plots.
    """
    if self.writer:
        for pf in plotting_functions:
            descr, fig = pf(model, iteration)
            self.writer.add_figure(descr, fig, global_step=iteration)
    else:
        raise RuntimeError(
            "The writer is not initialized."
            "Please call the 'writer.open()' method before writing"
        )

`write(iteration, metrics)`

Logs the results of the current iteration to TensorBoard.

PARAMETER	DESCRIPTION
`iteration`	The current training iteration. TYPE: `int`
`metrics`	A dictionary of metrics to log, where keys are metric names and values are the corresponding metric values. TYPE: `dict`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def write(self, iteration: int, metrics: dict) -> None:
    """
    Logs the results of the current iteration to TensorBoard.

    Args:
        iteration (int): The current training iteration.
        metrics (dict): A dictionary of metrics to log, where keys are metric names
                        and values are the corresponding metric values.
    """
    if self.writer:
        for key, value in metrics.items():
            self.writer.add_scalar(key, value, iteration)
    else:
        raise RuntimeError(
            "The writer is not initialized."
            "Please call the 'writer.open()' method before writing."
        )

`get_writer(tracking_tool)`

Factory method to get the appropriate writer based on the tracking tool.

PARAMETER	DESCRIPTION
`tracking_tool`	The experiment tracking tool to use. TYPE: `ExperimentTrackingTool`

RETURNS	DESCRIPTION
`BaseWriter`	An instance of the appropriate writer. TYPE: `BaseWriter`

Source code in qadence/ml_tools/callbacks/writer_registry.py

def get_writer(tracking_tool: ExperimentTrackingTool) -> BaseWriter:
    """Factory method to get the appropriate writer based on the tracking tool.

    Args:
        tracking_tool (ExperimentTrackingTool): The experiment tracking tool to use.

    Returns:
        BaseWriter: An instance of the appropriate writer.
    """
    writer_class = WRITER_REGISTRY.get(tracking_tool)
    if writer_class:
        return writer_class()
    else:
        raise ValueError(f"Unsupported tracking tool: {tracking_tool}")

`InformationContent(model, loss_fn, xs, epsilons, variation_multiple=20)`

Information Landscape class.

This class handles the study of loss landscape from information theoretic perspective and provides methods to get bounds on the norm of the gradient from the Information Content of the loss landscape.

PARAMETER	DESCRIPTION
`model`	The quantum or classical model to analyze. TYPE: `Module`
`loss_fn`	Loss function that takes model output and calculates loss TYPE: `Callable`
`xs`	Input data to evaluate the model on TYPE: `Any`
`epsilons`	The thresholds to use for discretization of the finite derivatives TYPE: `Tensor`
`variation_multiple`	The number of sets of variational parameters to generate per each variational parameter. The number of variational parameters required for the statistical analysis scales linearly with the amount of them present in the model. This is that linear factor. TYPE: `int` DEFAULT: `20`

Notes

This class provides flexibility in terms of what the model, the loss function, and the xs are. The only requirement is that the loss_fn takes the model and xs as arguments and returns the loss, and another dictionary of other metrics.

Thus, assumed structure: loss_fn(model, xs) -> (loss, metrics, ...)

Example: A Classifier

model = nn.Linear(10, 1)

def loss_fn(
    model: nn.Module,
    xs: tuple[torch.Tensor, torch.Tensor]
) -> tuple[torch.Tensor, dict[str, float]:
    criterion = nn.MSELoss()
    inputs, labels = xs
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    metrics = {"loss": loss.item()}
    return loss, metrics

xs = (torch.randn(10, 10), torch.randn(10, 1))

info_landscape = InfoLandscape(model, loss_fn, xs)

In this example, the model is a linear classifier, and the xs include both the inputs and the target labels. The logic for calculation of the loss from this lies entirely within the loss_fn function. This can then further be used to obtain the bounds on the average norm of the gradient of the loss function.

Example: A Physics Informed Neural Network

class PhysicsInformedNN(nn.Module):
    // <Initialization Logic>

    def forward(self, xs: dict[str, torch.Tensor]):
        return {
            "pde_residual": pde_residual(xs["pde"]),
            "boundary_condition": bc_term(xs["bc"]),
        }

def loss_fn(
    model: PhysicsInformedNN,
    xs: dict[str, torch.Tensor]
) -> tuple[torch.Tensor, dict[str, float]:
    pde_residual, bc_term = model(xs)
    loss = torch.mean(torch.sum(pde_residual**2, dim=1), dim=0)
        + torch.mean(torch.sum(bc_term**2, dim=1), dim=0)

    return loss, {"pde_residual": pde_residual, "bc_term": bc_term}

xs = {
    "pde": torch.linspace(0, 1, 10),
    "bc": torch.tensor([0.0]),
}

info_landscape = InfoLandscape(model, loss_fn, xs)

In this example, the model is a Physics Informed Neural Network, and the `xs`
are the inputs to the different residual components of the model. The logic
for calculation of the residuals lies within the PhysicsInformedNN class, and
the loss function is defined to calculate the loss that is to be optimized
from these residuals. This can then further be used to obtain the
bounds on the average norm of the gradient of the loss function.

The first value that the loss_fn returns is the loss value that is being optimized. The function is also expected to return other value(s), often the metrics that are used to calculate the loss. These values are ignored for the purpose of this class.

Source code in qadence/ml_tools/information/information_content.py

def __init__(
    self,
    model: nn.Module,
    loss_fn: Callable,
    xs: Any,
    epsilons: torch.Tensor,
    variation_multiple: int = 20,
) -> None:
    """Information Landscape class.

    This class handles the study of loss landscape from information theoretic
    perspective and provides methods to get bounds on the norm of the
    gradient from the Information Content of the loss landscape.

    Args:
        model: The quantum or classical model to analyze.
        loss_fn: Loss function that takes model output and calculates loss
        xs: Input data to evaluate the model on
        epsilons: The thresholds to use for discretization of the finite derivatives
        variation_multiple: The number of sets of variational parameters to generate per each
            variational parameter. The number of variational parameters required for the
            statistical analysis scales linearly with the amount of them present in the
            model. This is that linear factor.

    Notes:
        This class provides flexibility in terms of what the model, the loss function,
        and the xs are. The only requirement is that the loss_fn takes the model and xs as
        arguments and returns the loss, and another dictionary of other metrics.

        Thus, assumed structure:
            loss_fn(model, xs) -> (loss, metrics, ...)

        Example: A Classifier
            ```python
            model = nn.Linear(10, 1)

            def loss_fn(
                model: nn.Module,
                xs: tuple[torch.Tensor, torch.Tensor]
            ) -> tuple[torch.Tensor, dict[str, float]:
                criterion = nn.MSELoss()
                inputs, labels = xs
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                metrics = {"loss": loss.item()}
                return loss, metrics

            xs = (torch.randn(10, 10), torch.randn(10, 1))

            info_landscape = InfoLandscape(model, loss_fn, xs)
            ```
            In this example, the model is a linear classifier, and the `xs` include both the
            inputs and the target labels. The logic for calculation of the loss from this lies
            entirely within the `loss_fn` function. This can then further be used to obtain the
            bounds on the average norm of the gradient of the loss function.

        Example: A Physics Informed Neural Network
            ```python
            class PhysicsInformedNN(nn.Module):
                // <Initialization Logic>

                def forward(self, xs: dict[str, torch.Tensor]):
                    return {
                        "pde_residual": pde_residual(xs["pde"]),
                        "boundary_condition": bc_term(xs["bc"]),
                    }

            def loss_fn(
                model: PhysicsInformedNN,
                xs: dict[str, torch.Tensor]
            ) -> tuple[torch.Tensor, dict[str, float]:
                pde_residual, bc_term = model(xs)
                loss = torch.mean(torch.sum(pde_residual**2, dim=1), dim=0)
                    + torch.mean(torch.sum(bc_term**2, dim=1), dim=0)

                return loss, {"pde_residual": pde_residual, "bc_term": bc_term}

            xs = {
                "pde": torch.linspace(0, 1, 10),
                "bc": torch.tensor([0.0]),
            }

            info_landscape = InfoLandscape(model, loss_fn, xs)
            ```

            In this example, the model is a Physics Informed Neural Network, and the `xs`
            are the inputs to the different residual components of the model. The logic
            for calculation of the residuals lies within the PhysicsInformedNN class, and
            the loss function is defined to calculate the loss that is to be optimized
            from these residuals. This can then further be used to obtain the
            bounds on the average norm of the gradient of the loss function.

        The first value that the `loss_fn` returns is the loss value that is being optimized.
        The function is also expected to return other value(s), often the metrics that are
        used to calculate the loss. These values are ignored for the purpose of this class.
    """
    self.model = model
    self.loss_fn = loss_fn
    self.xs = xs
    self.epsilons = epsilons
    self.device = next(model.parameters()).device

    self.param_shapes = {}
    self.total_params = 0

    for name, param in model.named_parameters():
        self.param_shapes[name] = param.shape
        self.total_params += param.numel()
    self.n_variations = variation_multiple * self.total_params
    self.all_variations = torch.empty(
        (self.n_variations, self.total_params), device=self.device
    ).uniform_(0, 2 * torch.pi)

`calculate_IC` `cached` `property`

Calculate Information Content for multiple epsilon values.

Returns: Tensor of IC values for each epsilon [n_epsilons]

`batched_loss()`

Calculate loss for all parameter variations in a batched manner.

Returns: Tensor of loss values for each parameter variation

Source code in qadence/ml_tools/information/information_content.py

def batched_loss(self) -> torch.Tensor:
    """Calculate loss for all parameter variations in a batched manner.

    Returns: Tensor of loss values for each parameter variation
    """
    param_variations = self.reshape_param_variations()
    losses = torch.zeros(self.n_variations, device=self.device)

    for i in range(self.n_variations):
        params = {name: param[i] for name, param in param_variations.items()}
        current_model = lambda x: functional_call(self.model, params, (x,))
        losses[i] = self.loss_fn(current_model, self.xs)[0]

    return losses

`calculate_transition_probabilities_batch()`

Calculate transition probabilities for multiple epsilon values.

RETURNS	DESCRIPTION
`Tensor`	Tensor of shape [n_epsilons, 6] containing probabilities for each transition type
`Tensor`	Columns order: [+1to0, +1to-1, 0to+1, 0to-1, -1to0, -1to+1]

Source code in qadence/ml_tools/information/information_content.py

def calculate_transition_probabilities_batch(self) -> torch.Tensor:
    """
    Calculate transition probabilities for multiple epsilon values.

    Returns:
        Tensor of shape [n_epsilons, 6] containing probabilities for each transition type
        Columns order: [+1to0, +1to-1, 0to+1, 0to-1, -1to0, -1to+1]
    """
    discretized = self.discretize_derivatives()

    current = discretized[:, :-1]
    next_val = discretized[:, 1:]

    transitions = torch.stack(
        [
            ((current == 1) & (next_val == 0)).sum(dim=1),
            ((current == 1) & (next_val == -1)).sum(dim=1),
            ((current == 0) & (next_val == 1)).sum(dim=1),
            ((current == 0) & (next_val == -1)).sum(dim=1),
            ((current == -1) & (next_val == 0)).sum(dim=1),
            ((current == -1) & (next_val == 1)).sum(dim=1),
        ],
        dim=1,
    ).float()

    total_transitions = current.size(1)
    probabilities = transitions / total_transitions

    return probabilities

`discretize_derivatives()`

Convert finite derivatives into discrete values.

RETURNS	DESCRIPTION
`Tensor`	Tensor containing discretized derivatives with shape [n_epsilons, n_variations-2]
`Tensor`	Each row contains {-1, 0, 1} values for that epsilon

Source code in qadence/ml_tools/information/information_content.py

def discretize_derivatives(self) -> torch.Tensor:
    """
    Convert finite derivatives into discrete values.

    Returns:
        Tensor containing discretized derivatives with shape [n_epsilons, n_variations-2]
        Each row contains {-1, 0, 1} values for that epsilon
    """
    derivatives = self.randomized_finite_der()

    derivatives = derivatives.unsqueeze(0)
    epsilons = self.epsilons.unsqueeze(1)

    discretized = torch.zeros((len(epsilons), len(derivatives[0])), device=self.device)
    discretized[derivatives > epsilons] = 1
    discretized[derivatives < -epsilons] = -1

    return discretized

`get_grad_norm_bounds_max_IC()`

Compute the bounds on the average norm of the gradient.

RETURNS	DESCRIPTION
`tuple[float, float]`	tuple[Tensor, Tensor]: The lower and upper bounds.

Source code in qadence/ml_tools/information/information_content.py

def get_grad_norm_bounds_max_IC(self) -> tuple[float, float]:
    """
    Compute the bounds on the average norm of the gradient.

    Returns:
        tuple[Tensor, Tensor]: The lower and upper bounds.
    """
    max_IC, epsilon_m = self.max_IC()
    lower_bound = (
        epsilon_m
        * sqrt(self.total_params)
        / (NormalDist().inv_cdf(1 - 2 * self.q_value(max_IC)))
    )
    upper_bound = (
        epsilon_m
        * sqrt(self.total_params)
        / (NormalDist().inv_cdf(0.5 * (1 + 2 * self.q_value(max_IC))))
    )

    if max_IC < log(2, 6):
        logger.warning(
            "Warning: The maximum IC is less than the required value. The bounds may be"
            + " inaccurate."
        )

    return lower_bound, upper_bound

`get_grad_norm_bounds_sensitivity_IC(eta)`

Compute the bounds on the average norm of the gradient.

PARAMETER	DESCRIPTION
`eta`	The sensitivity IC. TYPE: `float`

RETURNS	DESCRIPTION
`Tensor`	The lower bound. TYPE: `float`

Source code in qadence/ml_tools/information/information_content.py

def get_grad_norm_bounds_sensitivity_IC(self, eta: float) -> float:
    """
    Compute the bounds on the average norm of the gradient.

    Args:
        eta (float): The sensitivity IC.

    Returns:
        Tensor: The lower bound.
    """
    epsilon_sensitivity = self.sensitivity_IC(eta)
    upper_bound = (
        epsilon_sensitivity * sqrt(self.total_params) / (NormalDist().inv_cdf(1 - 3 * eta / 2))
    )
    return upper_bound

`max_IC()`

Get the maximum Information Content and its corresponding epsilon.

Returns: Tuple of (maximum IC value, optimal epsilon)

Source code in qadence/ml_tools/information/information_content.py

def max_IC(self) -> tuple[float, float]:
    """
    Get the maximum Information Content and its corresponding epsilon.

    Returns: Tuple of (maximum IC value, optimal epsilon)
    """
    max_ic, max_idx = torch.max(self.calculate_IC, dim=0)
    max_epsilon = self.epsilons[max_idx]
    return max_ic.item(), max_epsilon.item()

`q_value(H_value)` `cached` `staticmethod`

Compute the q value.

q is the solution to the equation: H(x) = 4h(x) + 2h(1/2 - 2x)

It is the value of the probability of 4 of the 6 transitions such that the IC is the same as the IC of our system.

This quantity is useful in calculating the bounds on the norms of the gradients.

PARAMETER	DESCRIPTION
`H_value`	The information content. TYPE: `float`

RETURNS	DESCRIPTION
`float`	The q value TYPE: `float`

Source code in qadence/ml_tools/information/information_content.py

@staticmethod
@functools.lru_cache
def q_value(H_value: float) -> float:
    """
    Compute the q value.

    q is the solution to the equation:
    H(x) = 4h(x) + 2h(1/2 - 2x)

    It is the value of the probability of 4 of the 6 transitions such that
    the IC is the same as the IC of our system.

    This quantity is useful in calculating the bounds on the norms of the gradients.

    Args:
        H_value (float): The information content.

    Returns:
        float: The q value
    """

    x = torch.linspace(0.001, 0.16667, 10000)

    H = -4 * x * torch.log(x) / torch.log(torch.tensor(6)) - 2 * (0.5 - 2 * x) * torch.log(
        0.5 - 2 * x
    ) / torch.log(torch.tensor(6))
    err = torch.abs(H - H_value)
    idx = torch.argmin(err)
    return float(x[idx].item())

`randomized_finite_der()`

Calculate normalized finite difference of loss on doing random walk in the parameter space.

This serves as a proxy for the derivative of the loss with respect to parameters.

RETURNS	DESCRIPTION
`Tensor`	Tensor containing normalized finite differences (approximate directional derivatives)
`Tensor`	between consecutive points in the random walk. Shape: [n_variations - 1]

Source code in qadence/ml_tools/information/information_content.py

def randomized_finite_der(self) -> torch.Tensor:
    """
    Calculate normalized finite difference of loss on doing random walk in the parameter space.

    This serves as a proxy for the derivative of the loss with respect to parameters.

    Returns:
        Tensor containing normalized finite differences (approximate directional derivatives)
        between consecutive points in the random walk. Shape: [n_variations - 1]
    """
    losses = self.batched_loss()

    return (losses[1:] - losses[:-1]) / (
        torch.norm(self.all_variations[1:] - self.all_variations[:-1], dim=1) + 1e-8
    )

`reshape_param_variations()`

Reshape variations of the model's variational parameters.

RETURNS	DESCRIPTION
`dict[str, Tensor]`	Dictionary of parameter tensors, each with shape [n_variations, *param_shape]

Source code in qadence/ml_tools/information/information_content.py

def reshape_param_variations(self) -> dict[str, torch.Tensor]:
    """Reshape variations of the model's variational parameters.

    Returns:
        Dictionary of parameter tensors, each with shape [n_variations, *param_shape]
    """
    param_variations = {}
    start_idx = 0

    for name, shape in self.param_shapes.items():
        param_size = torch.prod(torch.tensor(shape)).item()
        param_variations[name] = self.all_variations[
            :, start_idx : start_idx + param_size
        ].view(self.n_variations, *shape)
        start_idx += param_size

    return param_variations

`sensitivity_IC(eta)`

Find the minimum value of epsilon such that the information content is less than eta.

PARAMETER	DESCRIPTION
`eta`	Threshold value, the sensitivity IC. TYPE: `float`

Returns: The epsilon value that gives IC that is less than the sensitivity IC.

Source code in qadence/ml_tools/information/information_content.py

def sensitivity_IC(self, eta: float) -> float:
    """
    Find the minimum value of epsilon such that the information content is less than eta.

    Args:
        eta: Threshold value, the sensitivity IC.

    Returns: The epsilon value that gives IC that is less than the sensitivity IC.
    """
    ic_values = self.calculate_IC
    mask = ic_values < eta
    epsilons = self.epsilons[mask]
    return float(epsilons.min().item())

`Accelerator(nprocs=1, compute_setup='auto', log_setup='cpu', backend='gloo', dtype=None)`

Bases: Distributor

A class for handling distributed training.

This class extends Distributor to manage distributed training using PyTorch's torch.distributed API. It supports spawning multiple processes and wrapping models with DistributedDataParallel (DDP) when required.

This class is provides head level method - distribute() - which wraps a function at a head process level, before launching nprocs processes as required. Furthermore, it provides processes level methods, such as prepare(), and prepare_batch() which can be run inside each process for correct movement and preparation of model, optimizers and datasets.

Inherited Attributes

nprocs (int): Number of processes to launch for distributed training. execution (BaseExecution): Detected execution instance for process launch (e.g., "torchrun","default"). execution_type (ExecutionType): Type of execution used. rank (int): Global rank of the process (to be set during environment setup). world_size (int): Total number of processes (to be set during environment setup). local_rank (int | None): Local rank on the node (to be set during environment setup). master_addr (str): Master node address (to be set during environment setup). master_port (str): Master node port (to be set during environment setup). node_rank (int): Rank of the node on the cluster setup.

There are three different indicators for number of processes executed.

1. self._config_nprocs: Number of processes specified by the user. Provided in the initilization of the Accelerator. (acc = Accelerator(nprocs = 2))
1. self.nprocs: Number of processes defined at the head level.
2. When accelerator is used to spawn processes (e.g., In case default, python execution), nprocs = _config_nprocs.
3. When an external elastic method is used to spawn processes (e.g., In case of torchrun), nprocs = 1. This is because the external launcher already spawns multiple processes, and the accelerator init is called from each process.
1. self.world_size: Number of processes actually executed.

Initializes the Accelerator class.

PARAMETER	DESCRIPTION
`nprocs`	Number of processes to launch. Default is 1. TYPE: `int` DEFAULT: `1`
`compute_setup`	Compute device setup; options are "auto" (default), "gpu", or "cpu". - "auto": Uses GPU if available, otherwise CPU. - "gpu": Forces GPU usage, raising an error if no CUDA device is available. - "cpu": Forces CPU usage. TYPE: `str` DEFAULT: `'auto'`
`log_setup`	Logging device setup; options are "auto", "cpu" (default). - "auto": Uses same device to log as used for computation. - "cpu": Forces CPU logging. TYPE: `str` DEFAULT: `'cpu'`
`backend`	The backend for distributed communication. Default is "gloo". TYPE: `str` DEFAULT: `'gloo'`
`dtype`	Data type for controlling numerical precision. Default is None. TYPE: `dtype \| None` DEFAULT: `None`

Source code in qadence/ml_tools/train_utils/accelerator.py

def __init__(
    self,
    nprocs: int = 1,
    compute_setup: str = "auto",
    log_setup: str = "cpu",
    backend: str = "gloo",
    dtype: torch_dtype | None = None,
) -> None:
    """
    Initializes the Accelerator class.

    Args:
        nprocs (int): Number of processes to launch. Default is 1.
        compute_setup (str): Compute device setup; options are "auto" (default), "gpu", or "cpu".
            - "auto": Uses GPU if available, otherwise CPU.
            - "gpu": Forces GPU usage, raising an error if no CUDA device is available.
            - "cpu": Forces CPU usage.
        log_setup (str): Logging device setup; options are "auto", "cpu" (default).
            - "auto": Uses same device to log as used for computation.
            - "cpu": Forces CPU logging.
        backend (str): The backend for distributed communication. Default is "gloo".
        dtype (torch.dtype | None): Data type for controlling numerical precision. Default is None.
    """
    super().__init__(nprocs, compute_setup, log_setup, backend, dtype)

    # Default values
    self.rank = 0
    self.local_rank = 0
    self.world_size = self.execution.get_world_size(0, self.nprocs)

`_prepare_data(dataloader)`

Adjusts DataLoader(s) for distributed training.

PARAMETER	DESCRIPTION
`dataloader`	The dataloader or dictionary of dataloaders to prepare. TYPE: `Union[DataLoader, DictDataLoader]`

RETURNS	DESCRIPTION
`DataLoader \| DictDataLoader`	Union[DataLoader, DictDataLoader]: The prepared dataloader(s) with the correct distributed sampling setup.

Source code in qadence/ml_tools/train_utils/accelerator.py

def _prepare_data(self, dataloader: DataLoader | DictDataLoader) -> DataLoader | DictDataLoader:
    """
    Adjusts DataLoader(s) for distributed training.

    Args:
        dataloader (Union[DataLoader, DictDataLoader]): The dataloader or dictionary of dataloaders to prepare.

    Returns:
        Union[DataLoader, DictDataLoader]: The prepared dataloader(s) with the correct distributed
                                        sampling setup.
    """
    if isinstance(dataloader, DictDataLoader):
        # If the input is a DictDataLoader, prepare each contained DataLoader.
        prepared_dataloaders = {
            key: self._prepare_dataloader(dl) for key, dl in dataloader.dataloaders.items()
        }
        return DictDataLoader(prepared_dataloaders)
    else:
        # For a single DataLoader, prepare it directly.
        return self._prepare_dataloader(dataloader)

`_prepare_dataloader(dataloader)`

Prepares a single DataLoader for distributed training.

When training in a distributed setting (i.e., when self.world_size > 1), data must be divided among multiple processes. This is achieved by creating a DistributedSampler that splits the dataset into distinct subsets for each process.

This method does the following: - If distributed training is enabled: - Checks if the dataset is not an instance of InfiniteTensorDataset. - If so, creates a DistributedSampler for the dataset using the total number of replicas (self.world_size) and the current process's rank (self.local_rank). - Otherwise (i.e., for infinite datasets), no sampler is set (sampler remains None). - Returns a new DataLoader configured with: - The same dataset and batch size as the original. - The distributed sampler (if applicable). - The number of workers and pin_memory settings retrieved from the original DataLoader. - If not in a distributed setting (i.e., self.world_size <= 1), returns the original DataLoader unmodified.

PARAMETER	DESCRIPTION
`dataloader`	The original DataLoader instance that loads the dataset. TYPE: `DataLoader`

RETURNS	DESCRIPTION
`DataLoader`	A new DataLoader prepared for distributed training if in a multi-process environment; otherwise, the original DataLoader is returned. TYPE: `DataLoader`

Source code in qadence/ml_tools/train_utils/accelerator.py

def _prepare_dataloader(self, dataloader: DataLoader) -> DataLoader:
    """
    Prepares a single DataLoader for distributed training.

    When training in a distributed setting (i.e., when `self.world_size > 1`), data must be
    divided among multiple processes. This is achieved by creating a
    DistributedSampler that splits the dataset into distinct subsets for each process.

    This method does the following:
    - If distributed training is enabled:
        - Checks if the dataset is not an instance of `InfiniteTensorDataset`.
            - If so, creates a `DistributedSampler` for the dataset using the total number
                of replicas (`self.world_size`) and the current process's rank (`self.local_rank`).
            - Otherwise (i.e., for infinite datasets), no sampler is set (sampler remains `None`).
        - Returns a new DataLoader configured with:
            - The same dataset and batch size as the original.
            - The distributed sampler (if applicable).
            - The number of workers and pin_memory settings retrieved from the original DataLoader.
    - If not in a distributed setting (i.e., `self.world_size <= 1`), returns the original DataLoader unmodified.

    Args:
        dataloader (DataLoader): The original DataLoader instance that loads the dataset.

    Returns:
        DataLoader: A new DataLoader prepared for distributed training if in a multi-process environment;
                    otherwise, the original DataLoader is returned.
    """
    if self.world_size > 1:
        if not isinstance(dataloader.dataset, InfiniteTensorDataset):
            # If the dataset is not an infinite dataset, create a DistributedSampler.
            sampler = DistributedSampler(
                dataloader.dataset, num_replicas=self.world_size, rank=self.local_rank
            )
        else:
            # For infinite datasets, we do not use a sampler since the dataset
            # is designed to loop indefinitely.
            sampler = None

        return DataLoader(
            dataloader.dataset,  # Use the same dataset as the original.
            batch_size=dataloader.batch_size,  # Maintain the same batch size.
            sampler=sampler,  # Use the created DistributedSampler (or None).
            num_workers=getattr(dataloader, "num_workers", 0),
            pin_memory=getattr(dataloader, "pin_memory", False),
        )
    return dataloader

`_prepare_model(model)`

Moves the model to the desired device and casts it to the specified dtype.

In a distributed setting, if more than one device is used (i.e., self.world_size > 1), the model is wrapped in DistributedDataParallel (DDP) to handle gradient synchronization across devices.

PARAMETER	DESCRIPTION
`model`	The PyTorch model to prepare. TYPE: `Module`

RETURNS	DESCRIPTION
`Module`	nn.Module: The model moved to the correct device (and wrapped in DDP if applicable).

Source code in qadence/ml_tools/train_utils/accelerator.py

def _prepare_model(self, model: nn.Module) -> nn.Module:
    """
    Moves the model to the desired device and casts it to the specified dtype.

    In a distributed setting, if more than one device is used (i.e., self.world_size > 1),
    the model is wrapped in DistributedDataParallel (DDP) to handle gradient synchronization
    across devices.

    Args:
        model (nn.Module): The PyTorch model to prepare.

    Returns:
        nn.Module: The model moved to the correct device (and wrapped in DDP if applicable).
    """
    model = model.to(device=self.execution.device, dtype=self.execution.dtype)

    # If using distributed training with more than one device:
    if self.world_size > 1:
        if self.execution.device.startswith("cuda"):
            # For GPU-based training: wrap the model with DDP and specify the local GPU.
            model = DDP(model, device_ids=[self.local_rank])
        else:
            # For CPU-based or other environments:
            if not self.local_rank:
                model = DDP(model)

    return model

`_prepare_optimizer(optimizer)`

Passes through the optimizer without modification.

PARAMETER	DESCRIPTION
`optimizer`	The optimizer to prepare. TYPE: `Optimizer`

RETURNS	DESCRIPTION
`Optimizer`	optim.Optimizer: The unmodified optimizer.

Source code in qadence/ml_tools/train_utils/accelerator.py

def _prepare_optimizer(self, optimizer: optim.Optimizer) -> optim.Optimizer:
    """
    Passes through the optimizer without modification.

    Args:
        optimizer (optim.Optimizer): The optimizer to prepare.

    Returns:
        optim.Optimizer: The unmodified optimizer.
    """
    # Optimizers are not device-specific in this context, so no action is needed.
    return optimizer

`_spawn_method(instance, method, args, kwargs)`

This method spawns the required numbers of processes.

if execution is default, it will spawn nproc processes across all nodes
if execution is otherwise, it will run a single process.

PARAMETER	DESCRIPTION
`instance`	The object (Trainer) that contains the method to execute. This object is expected to have an `accelerator` attribute with a `setup_process(rank)` method. This argument is optional, in case it is None, the fun will be called independently. TYPE: `object`
`method`	The function of the method on the instance to be executed. TYPE: `Callable`
`args`	Positional arguments to pass to the target method. TYPE: `tuple`
`kwargs`	Keyword arguments to pass to the target method. TYPE: `dict`

Source code in qadence/ml_tools/train_utils/accelerator.py

def _spawn_method(self, instance: Any, method: Callable, args: Any, kwargs: Any) -> None:
    """
    This method spawns the required numbers of processes.

    - if execution is `default`, it will spawn `nproc` processes across all nodes
    - if execution is `otherwise`, it will run a single process.

    Args:
        instance (object): The object (Trainer) that contains the method to execute.
                           This object is expected to have an `accelerator` attribute with a `setup_process(rank)` method.
                           This argument is optional, in case it is None, the fun will be called independently.
        method (Callable): The function of the method on the instance to be executed.
        args (tuple): Positional arguments to pass to the target method.
        kwargs (dict): Keyword arguments to pass to the target method.
    """

    if self.execution_type == ExecutionType.DEFAULT and self.world_size > 1:
        # Spawn multiple processes that will run the worker function.
        nprocs = self.nprocs
        if self.execution.num_nodes > 1:
            nprocs //= self.execution.num_nodes
        mp.spawn(
            self.worker,
            args=(instance, method, args, kwargs),
            nprocs=int(nprocs),
            join=True,
        )
    else:
        # In single process mode, call the worker with rank 0.
        self.worker(0, instance, method, args, kwargs)

`all_reduce_dict(d, op='mean')`

Performs an all-reduce operation on a dictionary of tensors, averaging values across all processes.

PARAMETER	DESCRIPTION
`d`	A dictionary where values are tensors to be reduced across processes. TYPE: `dict[str, Tensor]`
`op`	Operation method to all_reduce with. Available options include `sum`, `avg`, and `max`. Defaults to `avg` TYPE: `str` DEFAULT: `'mean'`

RETURNS	DESCRIPTION
`dict[str, Tensor]`	dict[str, torch.Tensor]: A dictionary with the reduced tensors, averaged over the world size.

Source code in qadence/ml_tools/train_utils/accelerator.py

def all_reduce_dict(
    self, d: dict[str, torch.Tensor], op: str = "mean"
) -> dict[str, torch.Tensor]:
    """
    Performs an all-reduce operation on a dictionary of tensors, averaging values across all processes.

    Args:
        d (dict[str, torch.Tensor]): A dictionary where values are tensors to be reduced across processes.
        op (str): Operation method to all_reduce with. Available options include `sum`, `avg`, and `max`.
                        Defaults to `avg`

    Returns:
        dict[str, torch.Tensor]: A dictionary with the reduced tensors, averaged over the world size.
    """
    if dist.is_initialized():
        world_size = dist.get_world_size()
        reduced: dict[str, torch.Tensor] = {}
        for key, tensor in d.items():
            if not isinstance(tensor, torch.Tensor):
                tensor = torch.tensor(
                    tensor, device=self.execution.device, dtype=self.execution.data_dtype
                )
            tensor = tensor.detach().clone()
            if op == "max":
                dist.all_reduce(tensor, op=dist.ReduceOp.MAX)
            elif op == "sum":
                dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
            else:
                dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
                tensor /= world_size
            reduced[key] = tensor
        return reduced
    else:
        return d

`broadcast(obj, src)`

Broadcasts an object from the source process to all processes.

On non-source processes, this value is ignored.

PARAMETER	DESCRIPTION
`obj`	The object to broadcast on the source process. TYPE: `Any`
`src`	The source process rank. TYPE: `int`

RETURNS	DESCRIPTION
`Any`	The broadcasted object from the source process. TYPE: `Any`

Source code in qadence/ml_tools/train_utils/accelerator.py

def broadcast(self, obj: Any, src: int) -> Any:
    """
    Broadcasts an object from the source process to all processes.

    On non-source processes, this value is ignored.

    Args:
        obj (Any): The object to broadcast on the source process.
        src (int): The source process rank.

    Returns:
        Any : The broadcasted object from the source process.
    """
    if dist.is_initialized():
        obj_list = [obj] if self.rank == src else [None]
        dist.broadcast_object_list(obj_list, src=src)
        return obj_list[0]
    else:
        return obj

`distribute(fun)`

Decorator to distribute the fit function across multiple processes.

This function is generic and can work with other methods as well. Weather it is bound or unbound.

When applied to a function (typically a fit function), this decorator will execute the function in a distributed fashion using torch.multiprocessing. The number of processes used is determined by self.nprocs, and if multiple nodes are involved (self.num_nodes > 1), the process count is adjusted accordingly. In single process mode (self.nporcs is 1), the function is executed directly in the current process.

After execution, the decorator returns the model stored in instance.model.

PARAMETER	DESCRIPTION
`fun`	The function to be decorated. This function usually implements a model fitting or training routine. TYPE: `callable`

RETURNS	DESCRIPTION
`callable`	The wrapped function. When called, it will execute in distributed mode (if configured) and return the value of `instance.model`. TYPE: `Callable`

Source code in qadence/ml_tools/train_utils/accelerator.py

def distribute(self, fun: Callable) -> Callable:
    """
    Decorator to distribute the fit function across multiple processes.

    This function is generic and can work with other methods as well.
    Weather it is bound or unbound.

    When applied to a function (typically a fit function), this decorator
    will execute the function in a distributed fashion using torch.multiprocessing.
    The number of processes used is determined by `self.nprocs`,
    and if multiple nodes are involved (`self.num_nodes > 1`), the process count is
    adjusted accordingly. In single process mode (`self.nporcs` is 1), the function
    is executed directly in the current process.

    After execution, the decorator returns the model stored in `instance.model`.

    Parameters:
        fun (callable): The function to be decorated. This function usually implements
                        a model fitting or training routine.

    Returns:
        callable: The wrapped function. When called, it will execute in distributed mode
                  (if configured) and return the value of `instance.model`.
    """

    @functools.wraps(fun)
    def wrapper(*args: Any, **kwargs: Any) -> Any:

        # Get the original picklable function
        # for the case of bound class method
        # as well as a function
        if self.is_class_method(fun, args):
            instance = args[0]
            method_name = fun.__name__
            method = getattr(instance, method_name)
            args = args[1:]
            self._spawn_method(instance, method, args, kwargs)
        else:
            instance = None
            # method_name = fun.__name__
            # module = inspect.getmodule(fun)
            # method = getattr(module, method_name) if module else fun
            self._spawn_method(instance, fun, args, kwargs)

        if instance and hasattr(instance, "accelerator"):
            instance.accelerator.finalize()
        else:
            self.finalize()

        # TODO: Return the original returns from fun
        # Currently it only returns the model and optimizer
        # similar to the fit method.
        try:
            return instance.model, instance.optimizer
        except Exception:
            return

    return wrapper

`is_class_method(fun, args)`

Determines if fun is a class method or a standalone function.

Frist argument of the args should be: - An object and has dict: making it a class - Has a method named fun: making it a class that has this method.

PARAMETER	DESCRIPTION
`fun`	The function being checked. TYPE: `Callable`
`args`	The arguments passed to the function. TYPE: `tuple`

RETURNS	DESCRIPTION
`bool`	True if `fun` is a class method, False otherwise. TYPE: `bool`

Source code in qadence/ml_tools/train_utils/accelerator.py

def is_class_method(self, fun: Callable, args: Any) -> bool:
    """
    Determines if `fun` is a class method or a standalone function.

    Frist argument of the args should be:
    - An object and has __dict__: making it a class
    - Has a method named fun: making it a class that has this method.

    Args:
        fun (Callable): The function being checked.
        args (tuple): The arguments passed to the function.

    Returns:
        bool: True if `fun` is a class method, False otherwise.
    """
    return (
        bool(args)
        and isinstance(args[0], object)
        and hasattr(args[0], "__dict__")
        and hasattr(args[0], fun.__name__)
    )

`prepare(*args)`

Prepares models, optimizers, and dataloaders for distributed training.

This method iterates over the provided objects and: - Moves models to the specified device (e.g., GPU or CPU) and casts them to the desired precision (specified by self.dtype). It then wraps models in DistributedDataParallel (DDP) if more than one device is used. - Passes through optimizers unchanged. - For dataloaders, it adjusts them to use a distributed sampler (if applicable) by calling a helper method. Note that only the sampler is prepared; moving the actual batch data to the device is handled separately during training. Please use the prepare_batch method to move the batch to correct device/dtype.

PARAMETER	DESCRIPTION
`*args`	A variable number of objects to be prepared. These can include: - PyTorch models (`nn.Module`) - Optimizers (`optim.Optimizer`) - DataLoaders (or a dictionary-like `DictDataLoader` of dataloaders) TYPE: `Any` DEFAULT: `()`

RETURNS	DESCRIPTION
`tuple[Any, ...]`	tuple[Any, ...]: A tuple containing the prepared objects, where each object has been modified as needed to support distributed training.

Source code in qadence/ml_tools/train_utils/accelerator.py

def prepare(self, *args: Any) -> tuple[Any, ...]:
    """
    Prepares models, optimizers, and dataloaders for distributed training.

    This method iterates over the provided objects and:
    - Moves models to the specified device (e.g., GPU or CPU) and casts them to the
        desired precision (specified by `self.dtype`). It then wraps models in
        DistributedDataParallel (DDP) if more than one device is used.
    - Passes through optimizers unchanged.
    - For dataloaders, it adjusts them to use a distributed sampler (if applicable)
        by calling a helper method. Note that only the sampler is prepared; moving the
        actual batch data to the device is handled separately during training.
        Please use the `prepare_batch` method to move the batch to correct device/dtype.

    Args:
        *args (Any): A variable number of objects to be prepared. These can include:
            - PyTorch models (`nn.Module`)
            - Optimizers (`optim.Optimizer`)
            - DataLoaders (or a dictionary-like `DictDataLoader` of dataloaders)

    Returns:
        tuple[Any, ...]: A tuple containing the prepared objects, where each object has been
                        modified as needed to support distributed training.
    """
    prepared: list = []
    for obj in args:
        if obj is None:
            prepared.append(None)
        elif isinstance(obj, nn.Module):
            prepared.append(self._prepare_model(obj))
        elif isinstance(obj, optim.Optimizer):
            prepared.append(self._prepare_optimizer(obj))
        elif isinstance(obj, (DataLoader, DictDataLoader)):
            prepared.append(self._prepare_data(obj))
        else:
            prepared.append(obj)
    return tuple(prepared)

`prepare_batch(batch)`

Moves a batch of data to the target device and casts it to the desired data dtype.

This method is typically called within the optimization step of your training loop. It supports various batch formats: - If the batch is a dictionary, each value is moved individually. - If the batch is a tuple or list, each element is processed and returned as a tuple. - Otherwise, the batch is processed directly.

PARAMETER	DESCRIPTION
`batch`	The batch of data to move to the device. This can be a dict, tuple, list, or any type compatible with `data_to_device`. TYPE: `Any`

RETURNS	DESCRIPTION
`Any`	The batch with all elements moved to `self.device` and cast to `self.data_dtype`. TYPE: `Any`

Source code in qadence/ml_tools/train_utils/accelerator.py

def prepare_batch(self, batch: dict | list | tuple | torch.Tensor | None) -> Any:
    """
    Moves a batch of data to the target device and casts it to the desired data dtype.

    This method is typically called within the optimization step of your training loop.
    It supports various batch formats:
        - If the batch is a dictionary, each value is moved individually.
        - If the batch is a tuple or list, each element is processed and returned as a tuple.
        - Otherwise, the batch is processed directly.

    Args:
        batch (Any): The batch of data to move to the device. This can be a dict, tuple, list,
                     or any type compatible with `data_to_device`.

    Returns:
        Any: The batch with all elements moved to `self.device` and cast to `self.data_dtype`.
    """
    if batch is None:
        return None

    if isinstance(batch, dict):
        return {
            key: data_to_device(
                value, device=self.execution.device, dtype=self.execution.data_dtype
            )
            for key, value in batch.items()
        }
    elif isinstance(batch, (tuple, list)):
        return tuple(
            data_to_device(x, device=self.execution.device, dtype=self.execution.data_dtype)
            for x in batch
        )
    elif isinstance(batch, torch.Tensor):
        return data_to_device(
            batch, device=self.execution.device, dtype=self.execution.data_dtype
        )
    return

`worker(rank, instance, fun, args, kwargs)`

Worker function to be executed in each spawned process.

This function is called in every subprocess created by torch.multiprocessing (via mp.spawn). It performs the following tasks: 1. Sets up the accelerator for the given process rank. This typically involves configuring the GPU or other hardware resources for distributed training. 2. If the retrieved method has been decorated (i.e. it has a 'wrapped' attribute), the original, unwrapped function is invoked with the given arguments. Otherwise, the method is called directly.

PARAMETER	DESCRIPTION
`rank`	The rank (or identifier) of the spawned process. TYPE: `int`
`instance`	The object (Trainer) that contains the method to execute. This object is expected to have an `accelerator` attribute with a `setup_process(rank)` method. This argument is optional, in case it is None, the fun will be called independently. TYPE: `object`
`fun`	The function of the method on the instance to be executed. TYPE: `Callable`
`args`	Positional arguments to pass to the target method. TYPE: `tuple`
`kwargs`	Keyword arguments to pass to the target method. TYPE: `dict`

Source code in qadence/ml_tools/train_utils/accelerator.py

def worker(self, rank: int, instance: Any, fun: Callable, args: tuple, kwargs: dict) -> None:
    """
    Worker function to be executed in each spawned process.

    This function is called in every subprocess created by torch.multiprocessing (via mp.spawn).
    It performs the following tasks:
      1. Sets up the accelerator for the given process rank. This typically involves configuring
         the GPU or other hardware resources for distributed training.
      2. If the retrieved method has been decorated (i.e. it has a '__wrapped__' attribute),
         the original, unwrapped function is invoked with the given arguments. Otherwise,
         the method is called directly.

    Args:
        rank (int): The rank (or identifier) of the spawned process.
        instance (object): The object (Trainer) that contains the method to execute.
                           This object is expected to have an `accelerator` attribute with a `setup_process(rank)` method.
                           This argument is optional, in case it is None, the fun will be called independently.
        fun (Callable): The function of the method on the instance to be executed.
        args (tuple): Positional arguments to pass to the target method.
        kwargs (dict): Keyword arguments to pass to the target method.
    """
    # Setup the accelerator for the given process rank (e.g., configuring GPU)
    if instance and instance.accelerator:
        instance.accelerator.setup_process(rank)
    else:
        self.setup_process(rank)

    if hasattr(fun, "__wrapped__"):
        # Explicitly get the original (unbound) method, passing in the instance.
        # We need to call the original method in case so that MP spawn does not
        # create multiple processes. (To Avoid infinite loop)
        fun = fun.__wrapped__  # Unwrap if decorated
        fun(instance, *args, **kwargs) if instance else fun(*args, **kwargs)
    else:
        fun(*args, **kwargs)

QML tools

ML Tools

Trainer(model, optimizer, config, loss_fn='mse', train_dataloader=None, val_dataloader=None, test_dataloader=None, optimize_step=optimize_step, max_batches=None)

_aggregate_result(result)

_batch_iter(dataloader, num_batches)

_fit_end()

_fit_setup()

_modify_batch_end_loss_metrics(loss_metrics)

_train()

_train_epochs(epoch_start, epoch_end, train_task=None, val_task=None)

build_optimize_result(result)

fit(train_dataloader=None, val_dataloader=None)

get_ic_grad_bounds(eta, epsilons, variation_multiple=20, dataloader=None)

run_test_batch(batch)

run_train_batch(batch)

run_training(dataloader)

run_val_batch(batch)

run_validation(dataloader)

stop_training()

test(test_dataloader=None)

AnsatzConfig(depth=1, ansatz_type=AnsatzType.HEA, ansatz_strategy=Strategy.DIGITAL, strategy_args=dict(), m_block_qubits=None, param_prefix='theta', tag=None) dataclass

ansatz_strategy = Strategy.DIGITAL class-attribute instance-attribute

ansatz_type = AnsatzType.HEA class-attribute instance-attribute

depth = 1 class-attribute instance-attribute

m_block_qubits = None class-attribute instance-attribute

param_prefix = 'theta' class-attribute instance-attribute

strategy_args = field(default_factory=dict) class-attribute instance-attribute

tag = None class-attribute instance-attribute

basis_set = BasisSet.FOURIER class-attribute instance-attribute

feature_map_strategy = Strategy.DIGITAL class-attribute instance-attribute

feature_range = None class-attribute instance-attribute

inputs = None class-attribute instance-attribute

multivariate_strategy = MultivariateStrategy.PARALLEL class-attribute instance-attribute

num_features = 0 class-attribute instance-attribute

num_repeats = 0 class-attribute instance-attribute

operation = None class-attribute instance-attribute

param_prefix = None class-attribute instance-attribute

reupload_scaling = ReuploadScaling.CONSTANT class-attribute instance-attribute

tag = None class-attribute instance-attribute

target_range = None class-attribute instance-attribute

_subfolders = field(default_factory=list) class-attribute instance-attribute

all_reduce_metrics = False class-attribute instance-attribute

backend = 'gloo' class-attribute instance-attribute

batch_size = 1 class-attribute instance-attribute

callbacks = field(default_factory=lambda: list()) class-attribute instance-attribute

checkpoint_best_only = False class-attribute instance-attribute

checkpoint_every = 0 class-attribute instance-attribute

compute_setup = 'cpu' class-attribute instance-attribute

create_subfolder_per_run = False class-attribute instance-attribute

dtype = None class-attribute instance-attribute

hyperparams = field(default_factory=dict) class-attribute instance-attribute

log_folder = Path('./') class-attribute instance-attribute

log_model = False class-attribute instance-attribute

log_setup = 'cpu' class-attribute instance-attribute

max_iter = 10000 class-attribute instance-attribute

nprocs = 1 class-attribute instance-attribute

plot_every = 0 class-attribute instance-attribute

plotting_functions = field(default_factory=tuple) class-attribute instance-attribute

print_every = 0 class-attribute instance-attribute

root_folder = Path('./qml_logs') class-attribute instance-attribute

tracking_tool = ExperimentTrackingTool.TENSORBOARD class-attribute instance-attribute

trainstop_criterion = None class-attribute instance-attribute

val_epsilon = 1e-05 class-attribute instance-attribute

val_every = 0 class-attribute instance-attribute

validation_criterion = None class-attribute instance-attribute

verbose = True class-attribute instance-attribute

write_every = 0 class-attribute instance-attribute

get_parameters(model)

num_parameters(model)

set_parameters(model, theta)

optimize_step(model, optimizer, loss_fn, xs, device=None, dtype=None)

update_ng_parameters(model, optimizer, loss_fn, data, ng_params)

DictDataLoader(dataloaders) dataclass

InfiniteTensorDataset(*tensors)

OptimizeResult(iteration, model, optimizer, loss=None, metrics=lambda: dict()(), extra=lambda: dict()(), rank=0, device='cpu') dataclass

device = 'cpu' class-attribute instance-attribute

extra = field(default_factory=lambda: dict()) class-attribute instance-attribute

iteration instance-attribute

loss = None class-attribute instance-attribute

metrics = field(default_factory=lambda: dict()) class-attribute instance-attribute

`Trainer(model, optimizer, config, loss_fn='mse', train_dataloader=None, val_dataloader=None, test_dataloader=None, optimize_step=optimize_step, max_batches=None)`

`_aggregate_result(result)`

`_batch_iter(dataloader, num_batches)`

`_fit_end()`

`_fit_setup()`

`_modify_batch_end_loss_metrics(loss_metrics)`

`_train()`

`_train_epochs(epoch_start, epoch_end, train_task=None, val_task=None)`

`build_optimize_result(result)`

`fit(train_dataloader=None, val_dataloader=None)`

`get_ic_grad_bounds(eta, epsilons, variation_multiple=20, dataloader=None)`

`run_test_batch(batch)`

`run_train_batch(batch)`

`run_training(dataloader)`

`run_val_batch(batch)`

`run_validation(dataloader)`

`stop_training()`

`test(test_dataloader=None)`

`AnsatzConfig(depth=1, ansatz_type=AnsatzType.HEA, ansatz_strategy=Strategy.DIGITAL, strategy_args=dict(), m_block_qubits=None, param_prefix='theta', tag=None)` `dataclass`

`ansatz_strategy = Strategy.DIGITAL` `class-attribute` `instance-attribute`

`ansatz_type = AnsatzType.HEA` `class-attribute` `instance-attribute`

`depth = 1` `class-attribute` `instance-attribute`

`m_block_qubits = None` `class-attribute` `instance-attribute`

`param_prefix = 'theta'` `class-attribute` `instance-attribute`

`strategy_args = field(default_factory=dict)` `class-attribute` `instance-attribute`

`tag = None` `class-attribute` `instance-attribute`

`basis_set = BasisSet.FOURIER` `class-attribute` `instance-attribute`

`feature_map_strategy = Strategy.DIGITAL` `class-attribute` `instance-attribute`

`feature_range = None` `class-attribute` `instance-attribute`

`inputs = None` `class-attribute` `instance-attribute`

`multivariate_strategy = MultivariateStrategy.PARALLEL` `class-attribute` `instance-attribute`

`num_features = 0` `class-attribute` `instance-attribute`

`num_repeats = 0` `class-attribute` `instance-attribute`

`operation = None` `class-attribute` `instance-attribute`

`param_prefix = None` `class-attribute` `instance-attribute`

`reupload_scaling = ReuploadScaling.CONSTANT` `class-attribute` `instance-attribute`

`tag = None` `class-attribute` `instance-attribute`

`target_range = None` `class-attribute` `instance-attribute`

`_subfolders = field(default_factory=list)` `class-attribute` `instance-attribute`

`all_reduce_metrics = False` `class-attribute` `instance-attribute`

`backend = 'gloo'` `class-attribute` `instance-attribute`

`batch_size = 1` `class-attribute` `instance-attribute`

`callbacks = field(default_factory=lambda: list())` `class-attribute` `instance-attribute`

`checkpoint_best_only = False` `class-attribute` `instance-attribute`

`checkpoint_every = 0` `class-attribute` `instance-attribute`

`compute_setup = 'cpu'` `class-attribute` `instance-attribute`

`create_subfolder_per_run = False` `class-attribute` `instance-attribute`

`dtype = None` `class-attribute` `instance-attribute`

`hyperparams = field(default_factory=dict)` `class-attribute` `instance-attribute`

`log_folder = Path('./')` `class-attribute` `instance-attribute`

`log_model = False` `class-attribute` `instance-attribute`

`log_setup = 'cpu'` `class-attribute` `instance-attribute`

`max_iter = 10000` `class-attribute` `instance-attribute`

`nprocs = 1` `class-attribute` `instance-attribute`

`plot_every = 0` `class-attribute` `instance-attribute`

`plotting_functions = field(default_factory=tuple)` `class-attribute` `instance-attribute`

`print_every = 0` `class-attribute` `instance-attribute`

`root_folder = Path('./qml_logs')` `class-attribute` `instance-attribute`

`tracking_tool = ExperimentTrackingTool.TENSORBOARD` `class-attribute` `instance-attribute`

`trainstop_criterion = None` `class-attribute` `instance-attribute`

`val_epsilon = 1e-05` `class-attribute` `instance-attribute`

`val_every = 0` `class-attribute` `instance-attribute`

`validation_criterion = None` `class-attribute` `instance-attribute`

`verbose = True` `class-attribute` `instance-attribute`

`write_every = 0` `class-attribute` `instance-attribute`

`get_parameters(model)`

`num_parameters(model)`

`set_parameters(model, theta)`

`optimize_step(model, optimizer, loss_fn, xs, device=None, dtype=None)`

`update_ng_parameters(model, optimizer, loss_fn, data, ng_params)`

`DictDataLoader(dataloaders)` `dataclass`

`InfiniteTensorDataset(*tensors)`

`OptimizeResult(iteration, model, optimizer, loss=None, metrics=lambda: dict()(), extra=lambda: dict()(), rank=0, device='cpu')` `dataclass`

`device = 'cpu'` `class-attribute` `instance-attribute`

`extra = field(default_factory=lambda: dict())` `class-attribute` `instance-attribute`

`iteration` `instance-attribute`

`loss = None` `class-attribute` `instance-attribute`

`metrics = field(default_factory=lambda: dict())` `class-attribute` `instance-attribute`

`model` `instance-attribute`

`optimizer` `instance-attribute`