qadence 1.7.2__py3-none-any.whl → 1.7.3__py3-none-any.whl

qadence/backends/pyqtorch/convert_ops.py

@@ -8,6 +8,7 @@ import pyqtorch as pyq
 import sympy
 import torch
 from pyqtorch.apply import apply_operator
+from pyqtorch.embed import Embedding
 from pyqtorch.matrices import _dagger
 from pyqtorch.time_dependent.sesolve import sesolve
 from pyqtorch.utils import is_diag
@@ -24,7 +25,7 @@ from torch import (
 )
 from torch import device as torch_device
 from torch import dtype as torch_dtype
-from torch.nn import Module
+from torch.nn import Module, ParameterDict
 
 from qadence.backends.utils import (
     finitediff,
@@ -200,7 +201,12 @@ class PyQHamiltonianEvolution(Module):
         elif isinstance(block.generator, Tensor):
             m = block.generator.to(dtype=cdouble)
             hmat = block_to_tensor(
-                MatrixBlock(m, qubit_support=block.qubit_support),
+                MatrixBlock(
+                    m,
+                    qubit_support=block.qubit_support,
+                    check_unitary=False,
+                    check_hermitian=True,
+                ),
                 qubit_support=self.qubit_support,
                 use_full_support=False,
             )
@@ -313,7 +319,8 @@ class PyQHamiltonianEvolution(Module):
     def forward(
         self,
         state: Tensor,
-        values: dict[str, Tensor],
+        values: dict[str, Tensor] | ParameterDict = dict(),
+        embedding: Embedding | None = None,
     ) -> Tensor:
         if getattr(self.block.generator, "is_time_dependent", False):  # type: ignore [union-attr]
 

qadence/blocks/matrix.py

@@ -60,7 +60,13 @@ class MatrixBlock(PrimitiveBlock):
     name = "MatrixBlock"
     matrix: torch.Tensor
 
-    def __init__(self, matrix: torch.Tensor | np.ndarray, qubit_support: tuple[int, ...]) -> None:
+    def __init__(
+        self,
+        matrix: torch.Tensor | np.ndarray,
+        qubit_support: tuple[int, ...],
+        check_unitary: bool = True,
+        check_hermitian: bool = False,
+    ) -> None:
         if isinstance(matrix, np.ndarray):
             matrix = torch.tensor(matrix)
         if matrix.ndim == 3 and matrix.size(0) == 1:
@@ -69,10 +75,12 @@ class MatrixBlock(PrimitiveBlock):
             raise TypeError("Please provide a 2D matrix.")
         if not self.is_square(matrix):
             raise TypeError("Please provide a square matrix.")
-        if not self.is_hermitian(matrix):
-            logger.warning("Provided matrix is not hermitian.")
-        if not self.is_unitary(matrix):
-            logger.warning("Provided matrix is not unitary.")
+        if check_hermitian:
+            if not self.is_hermitian(matrix):
+                logger.warning("Provided matrix is not hermitian.")
+        if check_unitary:
+            if not self.is_unitary(matrix):
+                logger.warning("Provided matrix is not unitary.")
         self.matrix = matrix.clone()
         super().__init__(qubit_support)
 

qadence/engines/torch/differentiable_expectation.py

@@ -151,8 +151,9 @@ class DifferentiableExpectation:
         return (
             AdjointExpectation.apply(
                 self.circuit.native,
-                self.observable[0].native,  # Currently, adjoint only supports a single observable.
                 self.state,
+                self.observable[0].native,  # Currently, adjoint only supports a single observable.
+                None,
                 self.param_values.keys(),
                 *self.param_values.values(),
             )

qadence/ml_tools/config.py

@@ -5,15 +5,25 @@ import os
 from dataclasses import dataclass, field, fields
 from logging import getLogger
 from pathlib import Path
-from typing import Callable, Optional, Type
+from typing import Callable, Type
+from uuid import uuid4
 
 from sympy import Basic
+from torch import Tensor
 
 from qadence.blocks.analog import AnalogBlock
 from qadence.blocks.primitive import ParametricBlock
 from qadence.operations import RX, AnalogRX
 from qadence.parameters import Parameter
-from qadence.types import AnsatzType, BasisSet, MultivariateStrategy, ReuploadScaling, Strategy
+from qadence.types import (
+    AnsatzType,
+    BasisSet,
+    ExperimentTrackingTool,
+    LoggablePlotFunction,
+    MultivariateStrategy,
+    ReuploadScaling,
+    Strategy,
+)
 
 logger = getLogger(__file__)
 
@@ -37,10 +47,14 @@ class TrainConfig:
     print_every: int = 1000
     """Print loss/metrics."""
     write_every: int = 50
-    """Write tensorboard logs."""
+    """Write loss and metrics with the tracking tool."""
     checkpoint_every: int = 5000
     """Write model/optimizer checkpoint."""
-    folder: Optional[Path] = None
+    plot_every: int = 5000
+    """Write figures."""
+    log_model: bool = False
+    """Logs a serialised version of the model."""
+    folder: Path | None = None
     """Checkpoint/tensorboard logs folder."""
     create_subfolder_per_run: bool = False
     """Checkpoint/tensorboard logs stored in subfolder with name `<timestamp>_<PID>`.
@@ -59,14 +73,38 @@ class TrainConfig:
 
     validation loss across previous iterations.
     """
-    validation_criterion: Optional[Callable] = None
+    validation_criterion: Callable | None = None
     """A boolean function which evaluates a given validation metric is satisfied."""
-    trainstop_criterion: Optional[Callable] = None
+    trainstop_criterion: Callable | None = None
     """A boolean function which evaluates a given training stopping metric is satisfied."""
     batch_size: int = 1
     """The batch_size to use when passing a list/tuple of torch.Tensors."""
     verbose: bool = True
     """Whether or not to print out metrics values during training."""
+    tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+    """The tracking tool of choice."""
+    hyperparams: dict = field(default_factory=dict)
+    """Hyperparameters to track."""
+    plotting_functions: tuple[LoggablePlotFunction, ...] = field(default_factory=tuple)  # type: ignore
+    """Functions for in-train plotting."""
+
+    # tensorboard only allows for certain types as hyperparameters
+    _tb_allowed_hyperparams_types: tuple = field(
+        default=(int, float, str, bool, Tensor), init=False, repr=False
+    )
+
+    def _filter_tb_hyperparams(self) -> None:
+        keys_to_remove = [
+            key
+            for key, value in self.hyperparams.items()
+            if not isinstance(value, TrainConfig._tb_allowed_hyperparams_types)
+        ]
+        if keys_to_remove:
+            logger.warning(
+                f"Tensorboard cannot log the following hyperparameters: {keys_to_remove}."
+            )
+            for key in keys_to_remove:
+                self.hyperparams.pop(key)
 
     def __post_init__(self) -> None:
         if self.folder:
@@ -81,6 +119,64 @@ class TrainConfig:
             self.trainstop_criterion = lambda x: x <= self.max_iter
         if self.validation_criterion is None:
             self.validation_criterion = lambda *x: False
+        if self.hyperparams and self.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+            self._filter_tb_hyperparams()
+        if self.tracking_tool == ExperimentTrackingTool.MLFLOW:
+            self._mlflow_config = MLFlowConfig()
+        if self.plotting_functions and self.tracking_tool != ExperimentTrackingTool.MLFLOW:
+            logger.warning("In-training plots are only available with mlflow tracking.")
+        if not self.plotting_functions and self.tracking_tool == ExperimentTrackingTool.MLFLOW:
+            logger.warning("Tracking with mlflow, but no plotting functions provided.")
+
+    @property
+    def mlflow_config(self) -> MLFlowConfig:
+        if self.tracking_tool == ExperimentTrackingTool.MLFLOW:
+            return self._mlflow_config
+        else:
+            raise AttributeError(
+                "mlflow_config is available only for with the mlflow tracking tool."
+            )
+
+
+class MLFlowConfig:
+    """
+    Configuration for mlflow tracking.
+
+    Example:
+
+        export MLFLOW_TRACKING_URI=tracking_uri
+        export MLFLOW_EXPERIMENT=experiment_name
+        export MLFLOW_RUN_NAME=run_name
+    """
+
+    def __init__(self) -> None:
+        import mlflow
+
+        self.tracking_uri: str = os.getenv("MLFLOW_TRACKING_URI", "")
+        """The URI of the mlflow tracking server.
+
+        An empty string, or a local file path, prefixed with file:/.
+        Data is stored locally at the provided file (or ./mlruns if empty).
+        """
+
+        self.experiment_name: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
+        """The name of the experiment.
+
+        If None or empty, a new experiment is created with a random UUID.
+        """
+
+        self.run_name: str = os.getenv("MLFLOW_RUN_NAME", str(uuid4()))
+        """The name of the run."""
+
+        mlflow.set_tracking_uri(self.tracking_uri)
+
+        # activate existing or create experiment
+        exp_filter_string = f"name = '{self.experiment_name}'"
+        if not mlflow.search_experiments(filter_string=exp_filter_string):
+            mlflow.create_experiment(name=self.experiment_name)
+
+        self.experiment = mlflow.set_experiment(self.experiment_name)
+        self.run = mlflow.start_run(run_name=self.run_name, nested=False)
 
 
 @dataclass

qadence/ml_tools/printing.py

@@ -1,7 +1,23 @@
 from __future__ import annotations
 
+from logging import getLogger
+from typing import Any, Callable, Union
+
+from matplotlib.figure import Figure
+from mlflow.models import infer_signature
+from torch import Tensor
+from torch.nn import Module
+from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
+from qadence.ml_tools.data import DictDataLoader
+from qadence.types import ExperimentTrackingTool
+
+logger = getLogger(__name__)
+
+PlottingFunction = Callable[[Module, int], tuple[str, Figure]]
+InputData = Union[Tensor, dict[str, Tensor]]
+
 
 def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:
     msg = " ".join(
@@ -20,5 +36,110 @@ def write_tensorboard(
         writer.add_scalar(key, arg, iteration)
 
 
-def log_hyperparams(writer: SummaryWriter, hyperparams: dict, metrics: dict) -> None:
+def log_hyperparams_tensorboard(writer: SummaryWriter, hyperparams: dict, metrics: dict) -> None:
     writer.add_hparams(hyperparams, metrics)
+
+
+def plot_tensorboard(
+    writer: SummaryWriter,
+    model: Module,
+    iteration: int,
+    plotting_functions: tuple[PlottingFunction],
+) -> None:
+    for pf in plotting_functions:
+        descr, fig = pf(model, iteration)
+        writer.add_figure(descr, fig, global_step=iteration)
+
+
+def log_model_tensorboard(
+    writer: SummaryWriter,
+    model: Module,
+    dataloader: Union[None, DataLoader, DictDataLoader],
+) -> None:
+    logger.warning("Model logging is not supported by tensorboard. No model will be logged.")
+
+
+def write_mlflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
+    writer.log_metrics({"loss": float(loss)}, step=iteration)  # type: ignore
+    writer.log_metrics(metrics, step=iteration)  # logs the single metrics
+
+
+def log_hyperparams_mlflow(writer: Any, hyperparams: dict, metrics: dict) -> None:
+    writer.log_params(hyperparams)  # type: ignore
+
+
+def plot_mlflow(
+    writer: Any,
+    model: Module,
+    iteration: int,
+    plotting_functions: tuple[PlottingFunction],
+) -> None:
+    for pf in plotting_functions:
+        descr, fig = pf(model, iteration)
+        writer.log_figure(fig, descr)
+
+
+def log_model_mlflow(
+    writer: Any, model: Module, dataloader: DataLoader | DictDataLoader | None
+) -> None:
+    if dataloader is not None:
+        xs: InputData
+        xs, *_ = next(iter(dataloader))
+        preds = model(xs)
+        if isinstance(xs, Tensor):
+            xs = xs.numpy()
+            preds = preds.detach().numpy()
+        elif isinstance(xs, dict):
+            for key, val in xs.items():
+                xs[key] = val.numpy()
+            for key, val in preds.items():
+                preds[key] = val.detach.numpy()
+        signature = infer_signature(xs, preds)
+    else:
+        signature = None
+    writer.pytorch.log_model(model, artifact_path="model", signature=signature)
+
+
+TRACKER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: write_tensorboard,
+    ExperimentTrackingTool.MLFLOW: write_mlflow,
+}
+
+LOGGER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: log_hyperparams_tensorboard,
+    ExperimentTrackingTool.MLFLOW: log_hyperparams_mlflow,
+}
+
+PLOTTER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: plot_tensorboard,
+    ExperimentTrackingTool.MLFLOW: plot_mlflow,
+}
+
+MODEL_LOGGER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: log_model_tensorboard,
+    ExperimentTrackingTool.MLFLOW: log_model_mlflow,
+}
+
+
+def write_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return TRACKER_MAPPING[tracking_tool](*args)
+
+
+def log_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return LOGGER_MAPPING[tracking_tool](*args)
+
+
+def plot_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return PLOTTER_MAPPING[tracking_tool](*args)
+
+
+def log_model_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return MODEL_LOGGER_MAPPING[tracking_tool](*args)

qadence/ml_tools/train_grad.py

@@ -1,10 +1,17 @@
 from __future__ import annotations
 
+import importlib
 import math
 from logging import getLogger
 from typing import Callable, Union
 
-from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from torch import complex128, float32, float64
 from torch import device as torch_device
 from torch import dtype as torch_dtype
@@ -16,8 +23,15 @@ from torch.utils.tensorboard import SummaryWriter
 from qadence.ml_tools.config import TrainConfig
 from qadence.ml_tools.data import DictDataLoader, data_to_device
 from qadence.ml_tools.optimize_step import optimize_step
-from qadence.ml_tools.printing import print_metrics, write_tensorboard
+from qadence.ml_tools.printing import (
+    log_model_tracker,
+    log_tracker,
+    plot_tracker,
+    print_metrics,
+    write_tracker,
+)
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
+from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
 
@@ -30,7 +44,6 @@ def train(
     loss_fn: Callable,
     device: torch_device = None,
     optimize_step: Callable = optimize_step,
-    write_tensorboard: Callable = write_tensorboard,
     dtype: torch_dtype = None,
 ) -> tuple[Module, Optimizer]:
     """Runs the training loop with gradient-based optimizer.
@@ -48,15 +61,11 @@ def train(
             the model
         optimizer: The optimizer to use.
         config: `TrainConfig` with additional training options.
-        loss_fn: Loss function returning (loss: float, metrics: dict[str, float])
+        loss_fn: Loss function returning (loss: float, metrics: dict[str, float], ...)
         device: String defining device to train on, pass 'cuda' for GPU.
         optimize_step: Customizable optimization callback which is called at every iteration.=
             The function must have the signature `optimize_step(model,
             optimizer, loss_fn, xs, device="cpu")`.
-        write_tensorboard: Customizable tensorboard logging callback which is
-            called every `config.write_every` iterations. The function must have
-            the signature `write_tensorboard(writer, loss, metrics, iteration)`
-            (see the example below).
         dtype: The dtype to use for the data.
 
     Example:
@@ -122,8 +131,11 @@ def train(
         model = model.module.to(device=device, dtype=dtype)
     else:
         model = model.to(device=device, dtype=dtype)
-    # initialize tensorboard
-    writer = SummaryWriter(config.folder, purge_step=init_iter)
+    # initialize tracking tool
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer = SummaryWriter(config.folder, purge_step=init_iter)
+    else:
+        writer = importlib.import_module("mlflow")
 
     perform_val = isinstance(config.val_every, int)
     if perform_val:
@@ -166,7 +178,7 @@ def train(
                 best_val_loss, metrics = loss_fn(model, xs_to_device)
 
                 metrics["val_loss"] = best_val_loss
-                write_tensorboard(writer, None, metrics, init_iter)
+                write_tracker(writer, None, metrics, init_iter, tracking_tool=config.tracking_tool)
 
             if config.folder:
                 if config.checkpoint_best_only:
@@ -174,6 +186,14 @@ def train(
                 else:
                     write_checkpoint(config.folder, model, optimizer, init_iter)
 
+            plot_tracker(
+                writer,
+                model,
+                init_iter,
+                config.plotting_functions,
+                tracking_tool=config.tracking_tool,
+            )
+
         except KeyboardInterrupt:
             logger.info("Terminating training gracefully after the current iteration.")
 
@@ -218,19 +238,31 @@ def train(
                     print_metrics(loss, metrics, iteration - 1)
 
                 if iteration % config.write_every == 0:
-                    write_tensorboard(writer, loss, metrics, iteration - 1)
+                    write_tracker(
+                        writer, loss, metrics, iteration, tracking_tool=config.tracking_tool
+                    )
 
+                if iteration % config.plot_every == 0:
+                    plot_tracker(
+                        writer,
+                        model,
+                        iteration,
+                        config.plotting_functions,
+                        tracking_tool=config.tracking_tool,
+                    )
                 if perform_val:
                     if iteration % config.val_every == 0:
                         xs = next(dl_iter_val)
                         xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
-                        val_loss, _ = loss_fn(model, xs_to_device)
+                        val_loss, *_ = loss_fn(model, xs_to_device)
                         if config.validation_criterion(val_loss, best_val_loss, config.val_epsilon):  # type: ignore[misc]
                             best_val_loss = val_loss
                             if config.folder and config.checkpoint_best_only:
                                 write_checkpoint(config.folder, model, optimizer, iteration="best")
                             metrics["val_loss"] = val_loss
-                            write_tensorboard(writer, None, metrics, iteration)
+                            write_tracker(
+                                writer, loss, metrics, iteration, tracking_tool=config.tracking_tool
+                            )
 
                 if config.folder:
                     if iteration % config.checkpoint_every == 0 and not config.checkpoint_best_only:
@@ -245,17 +277,32 @@ def train(
         try:
             xs = next(dl_iter) if dataloader is not None else None  # type: ignore[arg-type]
             xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
-            loss, metrics = loss_fn(model, xs_to_device)
+            loss, metrics, *_ = loss_fn(model, xs_to_device)
+            if dataloader is None:
+                loss = loss.item()
             if iteration % config.print_every == 0 and config.verbose:
                 print_metrics(loss, metrics, iteration)
 
         except KeyboardInterrupt:
             logger.info("Terminating training gracefully after the current iteration.")
 
-    # Final printing, writing and checkpointing
+    # Final checkpointing and writing
     if config.folder and not config.checkpoint_best_only:
         write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tensorboard(writer, loss, metrics, iteration)
-    writer.close()
+    write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
+
+    # writing hyperparameters
+    if config.hyperparams:
+        log_tracker(writer, config.hyperparams, metrics, tracking_tool=config.tracking_tool)
+
+    # logging the model
+    if config.log_model:
+        log_model_tracker(writer, model, dataloader, tracking_tool=config.tracking_tool)
+
+    # close tracker
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer.close()
+    elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:
+        writer.end_run()
 
     return model, optimizer

qadence/ml_tools/train_no_grad.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
+import importlib
 from logging import getLogger
 from typing import Callable
 
 import nevergrad as ng
 from nevergrad.optimization.base import Optimizer as NGOptimizer
-from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from torch import Tensor
 from torch.nn import Module
 from torch.utils.data import DataLoader
@@ -14,9 +21,16 @@ from torch.utils.tensorboard import SummaryWriter
 from qadence.ml_tools.config import TrainConfig
 from qadence.ml_tools.data import DictDataLoader
 from qadence.ml_tools.parameters import get_parameters, set_parameters
-from qadence.ml_tools.printing import print_metrics, write_tensorboard
+from qadence.ml_tools.printing import (
+    log_model_tracker,
+    log_tracker,
+    plot_tracker,
+    print_metrics,
+    write_tracker,
+)
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
 from qadence.ml_tools.tensors import promote_to_tensor
+from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
 
@@ -42,6 +56,7 @@ def train(
         dataloader: Dataloader constructed via `dictdataloader`
         optimizer: The optimizer to use taken from the Nevergrad library. If this is not
             the case the function will raise an AssertionError
+        config: `TrainConfig` with additional training options.
         loss_fn: Loss function returning (loss: float, metrics: dict[str, float])
     """
     init_iter = 0
@@ -63,8 +78,11 @@ def train(
     # TODO: support also Scipy optimizers
     assert isinstance(optimizer, NGOptimizer), "Use only optimizers from the Nevergrad library"
 
-    # initialize tensorboard
-    writer = SummaryWriter(config.folder, purge_step=init_iter)
+    # initialize tracking tool
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer = SummaryWriter(config.folder, purge_step=init_iter)
+    else:
+        writer = importlib.import_module("mlflow")
 
     # set optimizer configuration and initial parameters
     optimizer.budget = config.max_iter
@@ -100,7 +118,16 @@ def train(
                 print_metrics(loss, metrics, iteration)
 
             if iteration % config.write_every == 0:
-                write_tensorboard(writer, loss, metrics, iteration)
+                write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
+
+            if iteration % config.plot_every == 0:
+                plot_tracker(
+                    writer,
+                    model,
+                    iteration,
+                    config.plotting_functions,
+                    tracking_tool=config.tracking_tool,
+                )
 
             if config.folder:
                 if iteration % config.checkpoint_every == 0:
@@ -109,10 +136,22 @@ def train(
             if iteration >= init_iter + config.max_iter:
                 break
 
-    ## Final writing and stuff
+    # writing hyperparameters
+    if config.hyperparams:
+        log_tracker(writer, config.hyperparams, metrics, tracking_tool=config.tracking_tool)
+
+    if config.log_model:
+        log_model_tracker(writer, model, dataloader, tracking_tool=config.tracking_tool)
+
+    # Final writing and checkpointing
     if config.folder:
         write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tensorboard(writer, loss, metrics, iteration)
-    writer.close()
+    write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
+
+    # close tracker
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer.close()
+    elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:
+        writer.end_run()
 
     return model, optimizer

qadence/types.py

@@ -6,9 +6,11 @@ from typing import Callable, Iterable, Tuple, Union
 
 import numpy as np
 import sympy
+from matplotlib.figure import Figure
 from numpy.typing import ArrayLike
 from pyqtorch.utils import SolverType
 from torch import Tensor, pi
+from torch.nn import Module
 
 TNumber = Union[int, float, complex, np.int64, np.float64]
 """Union of python and numpy numeric types."""
@@ -445,3 +447,13 @@ class ObservableTransform:
     """Use the given values as min and max."""
     NONE = "none"
     """No transformation."""
+
+
+class ExperimentTrackingTool(StrEnum):
+    TENSORBOARD = "tensorboard"
+    """Use the tensorboard experiment tracker."""
+    MLFLOW = "mlflow"
+    """Use the ml-flow experiment tracker."""
+
+
+LoggablePlotFunction = Callable[[Module, int], tuple[str, Figure]]

{qadence-1.7.2.dist-info → qadence-1.7.3.dist-info}/METADATA

@@ -1,8 +1,8 @@
 Metadata-Version: 2.3
 Name: qadence
-Version: 1.7.2
+Version: 1.7.3
 Summary: Pasqal interface for circuit-based quantum computing SDKs
-Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>
+Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>, Giorgio Tosti Balducci <giorgio.tosti-balducci@pasqal.com>
 License: Apache 2.0
 License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
@@ -22,7 +22,7 @@ Requires-Dist: matplotlib
 Requires-Dist: nevergrad
 Requires-Dist: numpy
 Requires-Dist: openfermion
-Requires-Dist: pyqtorch==1.2.5
+Requires-Dist: pyqtorch==1.3.2
 Requires-Dist: pyyaml
 Requires-Dist: rich
 Requires-Dist: scipy
@@ -33,6 +33,7 @@ Requires-Dist: torch
 Provides-Extra: all
 Requires-Dist: braket; extra == 'all'
 Requires-Dist: libs; extra == 'all'
+Requires-Dist: mlflow; extra == 'all'
 Requires-Dist: protocols; extra == 'all'
 Requires-Dist: pulser; extra == 'all'
 Requires-Dist: visualization; extra == 'all'
@@ -51,6 +52,8 @@ Requires-Dist: optax; extra == 'horqrux'
 Requires-Dist: sympy2jax; extra == 'horqrux'
 Provides-Extra: libs
 Requires-Dist: qadence-libs; extra == 'libs'
+Provides-Extra: mlflow
+Requires-Dist: mlflow; extra == 'mlflow'
 Provides-Extra: protocols
 Requires-Dist: qadence-protocols; extra == 'protocols'
 Provides-Extra: pulser

{qadence-1.7.2.dist-info → qadence-1.7.3.dist-info}/RECORD

@@ -18,7 +18,7 @@ qadence/register.py,sha256=mwmvS6PcTY0F9cIhTUXG3NT73FIagfMCwVqYa4DrQrk,13001
 qadence/serial_expr_grammar.peg,sha256=z5ytL7do9kO8o4h-V5GrsDuLdso0KsRcMuIYURFfmAY,328
 qadence/serialization.py,sha256=qEET6Gu9u2aSibPve3bJrqDzK2_gO3RPDJjt4ZY8GbE,15596
 qadence/states.py,sha256=5QIOBBYs8e2uLFiMa8iMYZ-MvWIFEqkZAjNYx0SyYPI,14843
-qadence/types.py,sha256=6Dw8Ibtn0ZKRVg5DcW3O5129LIwYT-LLT2V3XjxFQek,10729
+qadence/types.py,sha256=Uex0set3DqJZb9Frkh9FvC5cp2ypp2U_A4e_y2XRDTo,11054
 qadence/utils.py,sha256=zb2j7wURfy8kazaS84r4t35vAeDpo4Tpm4HbmPH-kFA,9865
 qadence/analog/__init__.py,sha256=BCyS9R4KUjzUXN0Ax3b0eMo8ZAuSkGoJQVtZ4_pvAFs,279
 qadence/analog/addressing.py,sha256=fu5-xW9lquEbagApNp23S_ET1kl0iDtZUrIYSVNmw9s,6435
@@ -51,7 +51,7 @@ qadence/backends/pulser/waveforms.py,sha256=0uz95b7rUaUUtN0tuHBZmJ0H6UBmfHST_59o
 qadence/backends/pyqtorch/__init__.py,sha256=0OdVy6cq0oQggV48LO1WXdaZuSkDkz7OYNEPIkNAmfk,140
 qadence/backends/pyqtorch/backend.py,sha256=5ChkSD3D5totCMxwbzC31yIeOBxw6QEM5CS6qnq1Jqw,9287
 qadence/backends/pyqtorch/config.py,sha256=jK-if0OF6L_inP-oZhWI4-b8wcrOiK8-EVv3NYDOfBM,2056
-qadence/backends/pyqtorch/convert_ops.py,sha256=8xw43wWETIuD2N-5g-EZVeUYYOMnTWuIePfqqpQ3Mbg,14897
+qadence/backends/pyqtorch/convert_ops.py,sha256=lrc2_l-xsMY0fBAtbYsStB_TwIiG27Bur_tVLpaaNDA,15160
 qadence/blocks/__init__.py,sha256=H6jEA_CptkE-eoB4UfSbUiDszbxxhZwECV_TgoZWXoU,960
 qadence/blocks/abstract.py,sha256=QFwKPagbTrn3V4c2DHpBd-QL_mVIUXfbvyBLUdD6zw4,12023
 qadence/blocks/analog.py,sha256=ymnnlSVoW1XL05ZvnnHCqRTHuOXIEY_7E9M0PNKJZy4,10812
@@ -59,7 +59,7 @@ qadence/blocks/block_to_tensor.py,sha256=Sg7YGKUoPUUHKvyB8Khztrk7UYnV5SD451_3I00
 qadence/blocks/composite.py,sha256=z_lXRBVnh-DdvfZdv6T0ZEmVhlU76zBt72P_FGGa-PQ,8897
 qadence/blocks/embedding.py,sha256=XC3_U4Dqi9jvU1TVbl2bZQJAGNL8Sww89eKOVLCsfiQ,6752
 qadence/blocks/manipulate.py,sha256=kPmzej7mnWFoqTJA2CkGulT7hcPha0GGPARC8rjZltg,2387
-qadence/blocks/matrix.py,sha256=k6CC3zN2i6nd7_9S9u4fJAxy9wfkM287945GpArwOhY,3771
+qadence/blocks/matrix.py,sha256=r1JqyD3kzspLq3aYbdVCNbWpjUJXmUx2c34luIo8Mcc,3947
 qadence/blocks/primitive.py,sha256=RoEA9_VCI_8o4yg_pMe5T38z3LD6IFz9qlCiF3iHmOo,16631
 qadence/blocks/utils.py,sha256=iCJDi6HTYYaQQCoP3cdIKeCDuy8KQCxctrHN5QWXV-M,16349
 qadence/constructors/__init__.py,sha256=oEGuILUB8qEbSeaKV9Q-Tk-DAVx-U0wqn8VoSztVueg,984
@@ -90,7 +90,7 @@ qadence/engines/jax/differentiable_backend.py,sha256=W5rDA8wb-ECnFWoLj4dVugF9v1l
 qadence/engines/jax/differentiable_expectation.py,sha256=XBYHT1XKRuZfKxTcNy8KJpSDPt-2PR4ZCanImCPI9OI,3677
 qadence/engines/torch/__init__.py,sha256=iZFdD32ot0B0CVyC-f5dVViOBnqoalxa6M9Lj4WQuPE,160
 qadence/engines/torch/differentiable_backend.py,sha256=AWthwvKE8pCOih4dZ3tXxQX4W1ps9mBcvo7n4V9V24Y,3553
-qadence/engines/torch/differentiable_expectation.py,sha256=kAYl23Xq9MwkLm0MzmiIES-qUXc2zQsl9TuIF_c-qTE,9599
+qadence/engines/torch/differentiable_expectation.py,sha256=ojSsde_5PtpebYYBWYmU5Jj76m7pfDki2uqChWfepts,9621
 qadence/exceptions/__init__.py,sha256=BU6vWrI9mshzr1aTPm1Ticr_o_42GjTrWI4OZXhThsI,203
 qadence/exceptions/exceptions.py,sha256=4j_VJpx2sZ2Mir5BJUWu4nwb131FY1ygO4q8-XlyfRc,190
 qadence/measurements/__init__.py,sha256=RIjG9tVJMqhNzyj7maZI250Um0KgHl2PizDcKJag-JU,161
@@ -104,17 +104,17 @@ qadence/mitigations/analog_zne.py,sha256=g0QkjSdF-N9Dv2N8Oza4sylnjUMid5ea-4NCT9T
 qadence/mitigations/protocols.py,sha256=Jq9MyLujfTyWmc7XVUGYVRUkJT1MmZw-GgmWpVjmX2Y,1608
 qadence/mitigations/readout.py,sha256=HPfYmdjRlieUdOBMZTghFK4DRWfveM4KkDkEI0bMI0E,6262
 qadence/ml_tools/__init__.py,sha256=HP4xjldkUZ9_WbZEDgpl31qoP9st5SBbC-DjI5pkx3k,1054
-qadence/ml_tools/config.py,sha256=c3vvQiNXlNoJnOuMFfqAd5fVkmpa7EzCN_ztvPV1jBU,14152
+qadence/ml_tools/config.py,sha256=g-hFaVoG57p0elde0giSEP5_XTvyPquDg49xGOtj6gA,17686
 qadence/ml_tools/constructors.py,sha256=cE510DqCKBe4tImH90qHawEbXU-mlQuW9Wh15lUON6Q,27293
 qadence/ml_tools/data.py,sha256=8ZUFjhQSp94w7icX7RzM2J39Yo7P_T-AgjcThBc8miI,4283
 qadence/ml_tools/models.py,sha256=SjwAPbSl9zn9YqfmwqHc2lIXCkIpwG_ysz4jieRh7W0,16996
 qadence/ml_tools/optimize_step.py,sha256=ATXWmAqybJVK3QmAaDqVXB5mxjTo2MIi_e0a5WSPFpc,1800
 qadence/ml_tools/parameters.py,sha256=gew2Kq_5-RgRpaTvs8eauVhgo0sTqqDQEV6WHFEiLGM,1301
-qadence/ml_tools/printing.py,sha256=Mzdhmm-gPclhYL0NPN2zwJ19-kKQ4PDwFkefIOYEmzU,745
+qadence/ml_tools/printing.py,sha256=aNJdUdh6UkNFIvgOFVFNXpHc_ilJSezPGgl7it4o7Q4,4549
 qadence/ml_tools/saveload.py,sha256=jeYG7Y1ime0P06SMWOiCgWlci-xHdEPrAARfM-awDH8,5798
 qadence/ml_tools/tensors.py,sha256=xZ9ZRzOqEaMgLUGWQf1najDmL6iLuN1ojCGVFs1Tm94,1337
-qadence/ml_tools/train_grad.py,sha256=VcBNr7g9tfjnlyG7XwQVHtEivnki34LJZcgJuRxyHhs,10924
-qadence/ml_tools/train_no_grad.py,sha256=PrOfPwu6C-YqfFxnRkbeyOQzqSyjRrx4AZZd6C-1xRw,4705
+qadence/ml_tools/train_grad.py,sha256=tf162ZfK07NZeqmTDvA92kkojPxX8s2nwBP_VM2qSvw,12190
+qadence/ml_tools/train_no_grad.py,sha256=4AvJok882Vq6EVE1NNLicd0aww6Fku0laZm05dcILaU,5868
 qadence/ml_tools/utils.py,sha256=PW8FyoV0mG_DtN1U8njTDV5qxZ0EK4mnFwMAsLBArfk,1410
 qadence/noise/__init__.py,sha256=r0nR8uEZeB1M9pI2UisjWq0bjw50fPFfVGzIMev923g,147
 qadence/noise/protocols.py,sha256=-aZ06JvMnpxCeT5v5lI_RNPOLbb9Ju1Pi1AB6uAXxVE,1653
@@ -133,7 +133,7 @@ qadence/transpile/digitalize.py,sha256=iWRwYAYQsD2INHj0HNbGJriv_3fRCuBW1nDBrwtKS
 qadence/transpile/flatten.py,sha256=EdhSG5WyF56nbnxINNLqrHgY84MRM1YFjT3fR4aph5Q,3427
 qadence/transpile/invert.py,sha256=KAefHTG2AWr39aengVhXrzCtJPhrZC-ZnL6vYvmbnY0,4867
 qadence/transpile/transpile.py,sha256=6MRRkk1OS279L1fwUQjazA6qlfpbd-T_EJMKT8hAhOU,2721
-qadence-1.7.2.dist-info/METADATA,sha256=BUoNcgG5HnxJHM8FDiCpsvqqxwZ7J2by60wlR8b3oVI,9774
-qadence-1.7.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-qadence-1.7.2.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-qadence-1.7.2.dist-info/RECORD,,
+qadence-1.7.3.dist-info/METADATA,sha256=d3TUjDh_-Ragvuhicd9BWMr4orwEwNgjAXRnGWk_ruQ,9936
+qadence-1.7.3.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+qadence-1.7.3.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+qadence-1.7.3.dist-info/RECORD,,