From d0567f4eaedd85ad64ab04d7a499ab569d814f83 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Fri, 7 Jun 2024 16:34:48 +0200
Subject: [PATCH 01/51] Add write ml-flow placeholder class

---
 qadence/ml_tools/printing.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 333c8720f..b3703fc08 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -21,3 +21,7 @@ def write_tensorboard(
 
 def log_hyperparams(writer: SummaryWriter, hyperparams: dict, metrics: dict) -> None:
     writer.add_hparams(hyperparams, metrics)
+
+
+def write_mflow():
+    pass

From 9aa979dff68487a403391a7a8180f8a61f0187e6 Mon Sep 17 00:00:00 2001
From: seitzdom <dominik.seitz@pasqal.com>
Date: Tue, 11 Jun 2024 15:27:58 +0200
Subject: [PATCH 02/51] some prototype

---
 pyproject.toml                 |  1 +
 qadence/ml_tools/config.py     |  3 +++
 qadence/ml_tools/printing.py   | 20 +++++++++++++++++++-
 qadence/ml_tools/train_grad.py | 23 ++++++++++++++++-------
 qadence/types.py               |  7 +++++++
 5 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 697e14efb..6632a1588 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -74,6 +74,7 @@ horqrux = [
 protocols = ["qadence-protocols"]
 libs = ["qadence-libs"]
 dlprof = ["nvidia-pyindex", "nvidia-dlprof[pytorch]"]
+mlflow = ["mlflow"]
 all = [
   "pulser-core==0.18.0",
   "pulser-simulation==0.18.0",
diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index b62233d86..da566efee 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -6,6 +6,8 @@
 from pathlib import Path
 from typing import Callable, Optional
 
+from qadence.types import ExperimentTrackingTool
+
 
 @dataclass
 class TrainConfig:
@@ -46,6 +48,7 @@ class TrainConfig:
     """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
 
     def __post_init__(self) -> None:
         if self.folder:
diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index b3703fc08..22bb44c50 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -1,7 +1,11 @@
 from __future__ import annotations
 
+from typing import Any
+
 from torch.utils.tensorboard import SummaryWriter
 
+from qadence.types import ExperimentTrackingTool
+
 
 def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:
     msg = " ".join(
@@ -23,5 +27,19 @@ def log_hyperparams(writer: SummaryWriter, hyperparams: dict, metrics: dict) ->
     writer.add_hparams(hyperparams, metrics)
 
 
-def write_mflow():
+def write_mflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
+    # TODO for giorgio
+    # if we use the pytorch.autolog, we can just open a context
     pass
+
+
+TRACKER_MAPPING = {
+    ExperimentTrackingTool.TENSORBOARD: write_tensorboard,
+    ExperimentTrackingTool.MLFLOW: write_mflow,
+}
+
+
+def write_tracker(
+    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return TRACKER_MAPPING[tracking_tool](*args)
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index ee6fe0f2b..d5fbcaa6e 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib
 from logging import getLogger
 from typing import Callable, Union
 
@@ -15,8 +16,9 @@
 from qadence.ml_tools.config import TrainConfig
 from qadence.ml_tools.data import DictDataLoader
 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 print_metrics, write_tracker
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
+from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
 
@@ -29,7 +31,7 @@ def train(
     loss_fn: Callable,
     device: torch_device = None,
     optimize_step: Callable = optimize_step,
-    write_tensorboard: Callable = write_tensorboard,
+    write_tensorboard: Callable = write_tracker,
     dtype: torch_dtype = None,
 ) -> tuple[Module, Optimizer]:
     """Runs the training loop with gradient-based optimizer.
@@ -122,9 +124,16 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
         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("mflow")
+        with writer.run:
+            pass
+        writer.mlflow.pytorch.autolog(
+            log_every_n_step=config.write_every, log_models=False, log_datasets=False
+        )
     ## Training
     progress = Progress(
         TextColumn("[progress.description]{task.description}"),
@@ -175,7 +184,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                     print_metrics(loss, metrics, iteration)
 
                 if iteration % config.write_every == 0:
-                    write_tensorboard(writer, loss, metrics, iteration)
+                    write_tracker(writer, loss, metrics, iteration)
 
                 if config.folder:
                     if iteration % config.checkpoint_every == 0:
@@ -188,7 +197,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     # Final writing and checkpointing
     if config.folder:
         write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tensorboard(writer, loss, metrics, iteration)
+    write_tracker(writer, loss, metrics, iteration)
     writer.close()
 
     return model, optimizer
diff --git a/qadence/types.py b/qadence/types.py
index 07a6580a1..64f22352c 100644
--- a/qadence/types.py
+++ b/qadence/types.py
@@ -401,3 +401,10 @@ class ReadOutOptimization(StrEnum):
 
 ParamDictType = dict[str, ArrayLike]
 DifferentiableExpression = Callable[..., ArrayLike]
+
+
+class ExperimentTrackingTool(StrEnum):
+    TENSORBOARD = "tensorboard"
+    """Use the tensorboard experiment tracker."""
+    MLFLOW = "mlflow"
+    """Use the ml-flow experiment tracker."""

From e43bd92816c5ca974e2b5413b3d27e5f01952d42 Mon Sep 17 00:00:00 2001
From: seitzdom <dominik.seitz@pasqal.com>
Date: Tue, 11 Jun 2024 17:06:20 +0200
Subject: [PATCH 03/51] mlconfig stuff

---
 qadence/ml_tools/printing.py   |  4 ++--
 qadence/ml_tools/train_grad.py | 14 ++++++++------
 qadence/ml_tools/utils.py      | 27 +++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 22bb44c50..394f9264e 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -29,8 +29,8 @@ def log_hyperparams(writer: SummaryWriter, hyperparams: dict, metrics: dict) ->
 
 def write_mflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
     # TODO for giorgio
-    # if we use the pytorch.autolog, we can just open a context
-    pass
+    writer.log_params({"loss": loss}, {"iteration": iteration})
+    writer.log_params(metrics)
 
 
 TRACKER_MAPPING = {
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index d5fbcaa6e..0bbf8c385 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -18,6 +18,7 @@
 from qadence.ml_tools.optimize_step import optimize_step
 from qadence.ml_tools.printing import print_metrics, write_tracker
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
+from qadence.ml_tools.utils import MLFlowConfig
 from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
@@ -31,7 +32,7 @@ def train(
     loss_fn: Callable,
     device: torch_device = None,
     optimize_step: Callable = optimize_step,
-    write_tensorboard: Callable = write_tracker,
+    write_tracker: Callable = write_tracker,
     dtype: torch_dtype = None,
 ) -> tuple[Module, Optimizer]:
     """Runs the training loop with gradient-based optimizer.
@@ -128,12 +129,13 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer = SummaryWriter(config.folder, purge_step=init_iter)
     else:
+        mlflowconfig = MLFlowConfig()  # Set up credentials for mlflow tracking
         writer = importlib.import_module("mflow")
-        with writer.run:
-            pass
-        writer.mlflow.pytorch.autolog(
-            log_every_n_step=config.write_every, log_models=False, log_datasets=False
-        )
+        writer.set_experiment(mlflowconfig.EXPERIMENT)
+
+        # writer.mlflow.pytorch.autolog(
+        #     log_every_n_step=config.write_every, log_models=False, log_datasets=False
+        # )
     ## Training
     progress = Progress(
         TextColumn("[progress.description]{task.description}"),
diff --git a/qadence/ml_tools/utils.py b/qadence/ml_tools/utils.py
index 96db165bd..439157431 100644
--- a/qadence/ml_tools/utils.py
+++ b/qadence/ml_tools/utils.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import os
+from dataclasses import dataclass
 from functools import singledispatch
 from typing import Any
 
@@ -7,10 +9,13 @@
 
 from qadence.blocks import AbstractBlock, parameters
 from qadence.circuit import QuantumCircuit
+from qadence.logger import get_script_logger
 from qadence.ml_tools.models import TransformedModule
 from qadence.models import QNN, QuantumModel
 from qadence.parameters import Parameter, stringify
 
+logger = get_script_logger(__name__)
+
 
 @singledispatch
 def rand_featureparameters(
@@ -44,3 +49,25 @@ def _(qnn: QNN, batch_size: int = 1) -> dict[str, Tensor]:
 @rand_featureparameters.register
 def _(tm: TransformedModule, batch_size: int = 1) -> dict[str, Tensor]:
     return rand_featureparameters(tm.model, batch_size)
+
+
+@dataclass
+class MLFlowConfig:
+    """
+    Example:
+
+        export MLFLOW_TRACKING_URI=tracking_uri
+        export MLFLOW_TRACKING_USERNAME=username
+        export MLFLOW_TRACKING_PASSWORD=password
+    """
+
+    MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
+    MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
+    MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
+    EXPERIMENT: str = os.getenv("MLFLOW_EXPERIMENT", "")
+
+    def __post_init__(self) -> None:
+        if self.MLFLOW_TRACKING_USERNAME != "":
+            logger.info(
+                f"Intialized mlflow remote logging for user {self.MLFLOW_TRACKING_USERNAME}."
+            )

From 8d84ecdd9afa1ae8d15f5ce1f80879249d2cd513 Mon Sep 17 00:00:00 2001
From: seitzdom <dominik.seitz@pasqal.com>
Date: Thu, 13 Jun 2024 11:57:49 +0200
Subject: [PATCH 04/51] correct tracking

---
 qadence/ml_tools/train_grad.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index 0bbf8c385..bdcd0cbe1 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -186,7 +186,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                     print_metrics(loss, metrics, iteration)
 
                 if iteration % config.write_every == 0:
-                    write_tracker(writer, loss, metrics, iteration)
+                    write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
 
                 if config.folder:
                     if iteration % config.checkpoint_every == 0:
@@ -199,7 +199,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     # Final writing and checkpointing
     if config.folder:
         write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tracker(writer, loss, metrics, iteration)
+    write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
     writer.close()
 
     return model, optimizer

From 37657025532faee6f909a02682942107ea7e703a Mon Sep 17 00:00:00 2001
From: seitzdom <dominik.seitz@pasqal.com>
Date: Thu, 13 Jun 2024 12:37:05 +0200
Subject: [PATCH 05/51] working test

---
 qadence/ml_tools/printing.py         | 4 ++--
 qadence/ml_tools/train_grad.py       | 6 ++++--
 qadence/ml_tools/utils.py            | 3 ++-
 tests/ml_tools/test_checkpointing.py | 4 +++-
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 394f9264e..bb1926dd0 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -29,8 +29,8 @@ def log_hyperparams(writer: SummaryWriter, hyperparams: dict, metrics: dict) ->
 
 def write_mflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
     # TODO for giorgio
-    writer.log_params({"loss": loss}, {"iteration": iteration})
-    writer.log_params(metrics)
+    writer.log_metrics({"loss": float(loss), "iteration": iteration})  # type: ignore
+    writer.log_metrics(metrics)
 
 
 TRACKER_MAPPING = {
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index bdcd0cbe1..8e0dfcd89 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -130,8 +130,9 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
         writer = SummaryWriter(config.folder, purge_step=init_iter)
     else:
         mlflowconfig = MLFlowConfig()  # Set up credentials for mlflow tracking
-        writer = importlib.import_module("mflow")
+        writer = importlib.import_module("mlflow")
         writer.set_experiment(mlflowconfig.EXPERIMENT)
+        # writer = writer.start_run()
 
         # writer.mlflow.pytorch.autolog(
         #     log_every_n_step=config.write_every, log_models=False, log_datasets=False
@@ -200,6 +201,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     if config.folder:
         write_checkpoint(config.folder, model, optimizer, iteration)
     write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
-    writer.close()
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer.close()
 
     return model, optimizer
diff --git a/qadence/ml_tools/utils.py b/qadence/ml_tools/utils.py
index 439157431..ec734fee5 100644
--- a/qadence/ml_tools/utils.py
+++ b/qadence/ml_tools/utils.py
@@ -4,6 +4,7 @@
 from dataclasses import dataclass
 from functools import singledispatch
 from typing import Any
+from uuid import uuid4
 
 from torch import Tensor, rand
 
@@ -64,7 +65,7 @@ class MLFlowConfig:
     MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
     MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
     MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
-    EXPERIMENT: str = os.getenv("MLFLOW_EXPERIMENT", "")
+    EXPERIMENT: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
 
     def __post_init__(self) -> None:
         if self.MLFLOW_TRACKING_USERNAME != "":
diff --git a/tests/ml_tools/test_checkpointing.py b/tests/ml_tools/test_checkpointing.py
index 627dea7de..42f126a9e 100644
--- a/tests/ml_tools/test_checkpointing.py
+++ b/tests/ml_tools/test_checkpointing.py
@@ -143,7 +143,9 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         loss = criterion(out, torch.rand(1))
         return loss, {}
 
-    config = TrainConfig(folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1)
+    config = TrainConfig(
+        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool="mlflow"  # type: ignore
+    )
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
     ckpts = [tmp_path / Path(f"model_QNN_ckpt_00{i}_device_cpu.pt") for i in range(1, 9)]
     assert all(os.path.isfile(ckpt) for ckpt in ckpts)

From 6552b98068b08093d8109644e15811f684f96c82 Mon Sep 17 00:00:00 2001
From: seitzdom <dominik.seitz@pasqal.com>
Date: Thu, 13 Jun 2024 13:34:34 +0200
Subject: [PATCH 06/51] add mark

---
 tests/ml_tools/test_checkpointing.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/ml_tools/test_checkpointing.py b/tests/ml_tools/test_checkpointing.py
index 42f126a9e..61e0a96ad 100644
--- a/tests/ml_tools/test_checkpointing.py
+++ b/tests/ml_tools/test_checkpointing.py
@@ -4,6 +4,7 @@
 from itertools import count
 from pathlib import Path
 
+import pytest
 import torch
 from torch.utils.data import DataLoader
 
@@ -18,6 +19,7 @@
 from qadence.ml_tools.parameters import get_parameters, set_parameters
 from qadence.ml_tools.utils import rand_featureparameters
 from qadence.models import QNN, QuantumModel
+from qadence.types import ExperimentTrackingTool
 
 
 def dataloader(batch_size: int = 25) -> DataLoader:
@@ -129,7 +131,10 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     assert torch.allclose(loaded_model.expectation(inputs), model.expectation(inputs))
 
 
-def test_check_QNN_ckpts_exist(BasicQNN: QNN, tmp_path: Path) -> None:
+@pytest.mark.parametrize(
+    "tool", [ExperimentTrackingTool.TENSORBOARD, ExperimentTrackingTool.MLFLOW]
+)
+def test_check_QNN_ckpts_exist(tool: ExperimentTrackingTool, BasicQNN: QNN, tmp_path: Path) -> None:
     data = dataloader()
     model = BasicQNN
     cnt = count()
@@ -144,7 +149,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         return loss, {}
 
     config = TrainConfig(
-        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool="mlflow"  # type: ignore
+        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=tool
     )
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
     ckpts = [tmp_path / Path(f"model_QNN_ckpt_00{i}_device_cpu.pt") for i in range(1, 9)]

From ebda5ba8a66271aab0d012592152da72b4a4aadb Mon Sep 17 00:00:00 2001
From: seitzdom <dominik.seitz@pasqal.com>
Date: Thu, 13 Jun 2024 14:32:11 +0200
Subject: [PATCH 07/51] Working version of mlflow logging + demonstration

---
 docs/tutorials/qml/ml_tools.md             | 57 +++++++++++++++++++++
 docs/tutorials/qml/mlflow_demonstration.py | 58 ++++++++++++++++++++++
 qadence/ml_tools/config.py                 | 34 +++++++++++++
 qadence/ml_tools/printing.py               |  8 +--
 qadence/ml_tools/train_grad.py             |  9 ++--
 qadence/ml_tools/utils.py                  | 25 ----------
 6 files changed, 157 insertions(+), 34 deletions(-)
 create mode 100644 docs/tutorials/qml/mlflow_demonstration.py

diff --git a/docs/tutorials/qml/ml_tools.md b/docs/tutorials/qml/ml_tools.md
index 92921bc03..c5f40fdaf 100644
--- a/docs/tutorials/qml/ml_tools.md
+++ b/docs/tutorials/qml/ml_tools.md
@@ -284,3 +284,60 @@ def train(
 
     return model, optimizer
 ```
+
+## MLTools with mlflow
+
+MLTools now offers MLflow support.
+
+
+```python
+from __future__ import annotations
+
+import os
+from itertools import count
+import torch
+from torch.utils.data import DataLoader
+
+from qadence.ml_tools import (
+    TrainConfig,
+    train_with_grad,
+
+)
+from qadence.ml_tools.data import to_dataloader
+from qadence.ml_tools.utils import rand_featureparameters
+from qadence.models import QNN, QuantumModel
+from qadence.types import ExperimentTrackingTool
+from qadence import QuantumCircuit, hea, Z
+
+os.environ['MLFLOW_TRACKING_URI'] = 'sqlite:///mlflow.db'
+os.environ['MLFLOW_EXPERIMENT'] = 'mlflow_demonstration'
+os.environ['MLFLOW_RUN_NAME'] = 'test_0'
+
+# in case you want to track remotely
+#os.environ['MLFLOW_TRACKING_USERNAME'] =
+#s.environ['MLFLOW_TRACKING_PASSWORD'] =
+def dataloader(batch_size: int = 25) -> DataLoader:
+    x = torch.linspace(0, 1, batch_size).reshape(-1, 1)
+    y = torch.cos(x)
+    return to_dataloader(x, y, batch_size=batch_size, infinite=True)
+data = dataloader()
+model = QNN(QuantumCircuit(2, hea(2,1)), observable=Z(0))
+cnt = count()
+criterion = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+inputs = rand_featureparameters(model, 1)
+
+def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+    next(cnt)
+    out = model.expectation(inputs)
+    loss = criterion(out, torch.rand(1))
+    return loss, {}
+
+config = TrainConfig(
+    folder='mlflow_demonstration', max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=ExperimentTrackingTool.MLFLOW
+)
+train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+os.system('mlflow ui --port 5000')
+os.system('mlflow ui --backend-store-uri sqlite:///mlflow.db')
+```
diff --git a/docs/tutorials/qml/mlflow_demonstration.py b/docs/tutorials/qml/mlflow_demonstration.py
new file mode 100644
index 000000000..64fd23054
--- /dev/null
+++ b/docs/tutorials/qml/mlflow_demonstration.py
@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import os
+from itertools import count
+
+import torch
+from torch.utils.data import DataLoader
+
+from qadence import QuantumCircuit, Z, hea
+from qadence.ml_tools import (
+    TrainConfig,
+    train_with_grad,
+)
+from qadence.ml_tools.data import to_dataloader
+from qadence.ml_tools.utils import rand_featureparameters
+from qadence.models import QNN, QuantumModel
+from qadence.types import ExperimentTrackingTool
+
+os.environ["MLFLOW_TRACKING_URI"] = "sqlite:///mlflow.db"
+os.environ["MLFLOW_EXPERIMENT"] = "mlflow_demonstration"
+os.environ["MLFLOW_RUN_NAME"] = "test_0"
+
+
+# in case you want to track remotely
+# os.environ['MLFLOW_TRACKING_USERNAME'] =
+# s.environ['MLFLOW_TRACKING_PASSWORD'] =
+def dataloader(batch_size: int = 25) -> DataLoader:
+    x = torch.linspace(0, 1, batch_size).reshape(-1, 1)
+    y = torch.cos(x)
+    return to_dataloader(x, y, batch_size=batch_size, infinite=True)
+
+
+data = dataloader()
+model = QNN(QuantumCircuit(2, hea(2, 1)), observable=Z(0))
+cnt = count()
+criterion = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+inputs = rand_featureparameters(model, 1)
+
+
+def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+    next(cnt)
+    out = model.expectation(inputs)
+    loss = criterion(out, torch.rand(1))
+    return loss, {}
+
+
+config = TrainConfig(
+    folder="mlflow_demonstration",
+    max_iter=10,
+    checkpoint_every=1,
+    write_every=1,
+    tracking_tool=ExperimentTrackingTool.MLFLOW,
+)
+train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+os.system("mlflow ui --port 5000")
+os.system("mlflow ui --backend-store-uri sqlite:///mlflow.db")
diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index da566efee..39f161dc2 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -3,11 +3,15 @@
 import datetime
 import os
 from dataclasses import dataclass
+from logging import getLogger
 from pathlib import Path
 from typing import Callable, Optional
+from uuid import uuid4
 
 from qadence.types import ExperimentTrackingTool
 
+logger = getLogger(__name__)
+
 
 @dataclass
 class TrainConfig:
@@ -50,6 +54,8 @@ class TrainConfig:
     """Whether or not to print out metrics values during training."""
     tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
 
+    # mlflow_callbacks: list[Callable] = [write_mlflow_figure(), write_x()]
+
     def __post_init__(self) -> None:
         if self.folder:
             if isinstance(self.folder, str):  # type: ignore [unreachable]
@@ -63,3 +69,31 @@ def __post_init__(self) -> None:
             self.trainstop_criterion = lambda x: x <= self.max_iter
         if self.validation_criterion is None:
             self.validation_criterion = lambda x: False
+
+
+@dataclass
+class MLFlowConfig:
+    """
+    Example:
+
+        export MLFLOW_TRACKING_URI=tracking_uri
+        export MLFLOW_TRACKING_USERNAME=username
+        export MLFLOW_TRACKING_PASSWORD=password
+    """
+
+    MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
+    MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
+    MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
+    EXPERIMENT: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
+    RUN_NAME: str = os.getenv("MLFLOW_RUN_NAME", "test_0")
+
+    def __post_init__(self) -> None:
+        import mlflow
+
+        if self.MLFLOW_TRACKING_USERNAME != "":
+            logger.info(
+                f"Intialized mlflow remote logging for user {self.MLFLOW_TRACKING_USERNAME}."
+            )
+        mlflow.set_tracking_uri(self.MLFLOW_TRACKING_URI)
+        mlflow.set_experiment(self.EXPERIMENT)
+        mlflow.start_run(run_name=self.RUN_NAME, nested=False)
diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index bb1926dd0..11aba23db 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -27,15 +27,15 @@ def log_hyperparams(writer: SummaryWriter, hyperparams: dict, metrics: dict) ->
     writer.add_hparams(hyperparams, metrics)
 
 
-def write_mflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
+def write_mlflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
     # TODO for giorgio
-    writer.log_metrics({"loss": float(loss), "iteration": iteration})  # type: ignore
-    writer.log_metrics(metrics)
+    writer.log_metrics({"loss": float(loss)}, step=iteration)  # type: ignore
+    writer.log_metrics(metrics, step=iteration)
 
 
 TRACKER_MAPPING = {
     ExperimentTrackingTool.TENSORBOARD: write_tensorboard,
-    ExperimentTrackingTool.MLFLOW: write_mflow,
+    ExperimentTrackingTool.MLFLOW: write_mlflow,
 }
 
 
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index 8e0dfcd89..fdca919ac 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -13,12 +13,11 @@
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
-from qadence.ml_tools.config import TrainConfig
+from qadence.ml_tools.config import MLFlowConfig, TrainConfig
 from qadence.ml_tools.data import DictDataLoader
 from qadence.ml_tools.optimize_step import optimize_step
 from qadence.ml_tools.printing import print_metrics, write_tracker
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
-from qadence.ml_tools.utils import MLFlowConfig
 from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
@@ -129,10 +128,8 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer = SummaryWriter(config.folder, purge_step=init_iter)
     else:
-        mlflowconfig = MLFlowConfig()  # Set up credentials for mlflow tracking
+        MLFlowConfig()  # Set up credentials for mlflow tracking
         writer = importlib.import_module("mlflow")
-        writer.set_experiment(mlflowconfig.EXPERIMENT)
-        # writer = writer.start_run()
 
         # writer.mlflow.pytorch.autolog(
         #     log_every_n_step=config.write_every, log_models=False, log_datasets=False
@@ -203,5 +200,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer.close()
+    elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:
+        writer.end_run()
 
     return model, optimizer
diff --git a/qadence/ml_tools/utils.py b/qadence/ml_tools/utils.py
index ec734fee5..8006f5746 100644
--- a/qadence/ml_tools/utils.py
+++ b/qadence/ml_tools/utils.py
@@ -1,10 +1,7 @@
 from __future__ import annotations
 
-import os
-from dataclasses import dataclass
 from functools import singledispatch
 from typing import Any
-from uuid import uuid4
 
 from torch import Tensor, rand
 
@@ -50,25 +47,3 @@ def _(qnn: QNN, batch_size: int = 1) -> dict[str, Tensor]:
 @rand_featureparameters.register
 def _(tm: TransformedModule, batch_size: int = 1) -> dict[str, Tensor]:
     return rand_featureparameters(tm.model, batch_size)
-
-
-@dataclass
-class MLFlowConfig:
-    """
-    Example:
-
-        export MLFLOW_TRACKING_URI=tracking_uri
-        export MLFLOW_TRACKING_USERNAME=username
-        export MLFLOW_TRACKING_PASSWORD=password
-    """
-
-    MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
-    MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
-    MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
-    EXPERIMENT: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
-
-    def __post_init__(self) -> None:
-        if self.MLFLOW_TRACKING_USERNAME != "":
-            logger.info(
-                f"Intialized mlflow remote logging for user {self.MLFLOW_TRACKING_USERNAME}."
-            )

From 5198b49bbd03f1add165a526ee652a1fe69c6322 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 19 Jun 2024 15:07:37 +0200
Subject: [PATCH 08/51] Ignore mlflow runs

---
 .gitignore                                 | 3 +++
 docs/tutorials/qml/mlflow_demonstration.py | 7 ++-----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 1b8641509..a18f2e0e4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -75,3 +75,6 @@ events.out.tfevents.*
 *.dvi
 
 *.gv
+
+# mlflow
+mlruns/
diff --git a/docs/tutorials/qml/mlflow_demonstration.py b/docs/tutorials/qml/mlflow_demonstration.py
index 64fd23054..af16f2f87 100644
--- a/docs/tutorials/qml/mlflow_demonstration.py
+++ b/docs/tutorials/qml/mlflow_demonstration.py
@@ -6,11 +6,8 @@
 import torch
 from torch.utils.data import DataLoader
 
-from qadence import QuantumCircuit, Z, hea
-from qadence.ml_tools import (
-    TrainConfig,
-    train_with_grad,
-)
+from qadence import hea, QuantumCircuit, Z
+from qadence.ml_tools import train_with_grad, TrainConfig
 from qadence.ml_tools.data import to_dataloader
 from qadence.ml_tools.utils import rand_featureparameters
 from qadence.models import QNN, QuantumModel

From e5caf08400d8f86ba6cc5c7900549f46c599df92 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 19 Jun 2024 17:42:08 +0200
Subject: [PATCH 09/51] Add hyperparameter logging with mlflow

---
 docs/tutorials/qml/mlflow_demonstration.py | 21 +++++++++++++++++----
 qadence/ml_tools/config.py                 |  3 +++
 qadence/ml_tools/printing.py               | 20 +++++++++++++++++---
 qadence/ml_tools/train_grad.py             | 17 ++++++++++++++---
 4 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/docs/tutorials/qml/mlflow_demonstration.py b/docs/tutorials/qml/mlflow_demonstration.py
index af16f2f87..ab7cc2459 100644
--- a/docs/tutorials/qml/mlflow_demonstration.py
+++ b/docs/tutorials/qml/mlflow_demonstration.py
@@ -6,8 +6,8 @@
 import torch
 from torch.utils.data import DataLoader
 
-from qadence import hea, QuantumCircuit, Z
-from qadence.ml_tools import train_with_grad, TrainConfig
+from qadence import QuantumCircuit, Z, hea
+from qadence.ml_tools import TrainConfig, train_with_grad
 from qadence.ml_tools.data import to_dataloader
 from qadence.ml_tools.utils import rand_featureparameters
 from qadence.models import QNN, QuantumModel
@@ -17,6 +17,13 @@
 os.environ["MLFLOW_EXPERIMENT"] = "mlflow_demonstration"
 os.environ["MLFLOW_RUN_NAME"] = "test_0"
 
+hyperparams = {
+    "batch_size": 10,
+    "n_qubits": 2,
+    "ansatz_depth": 1,
+    "observable": Z,
+}
+
 
 # in case you want to track remotely
 # os.environ['MLFLOW_TRACKING_USERNAME'] =
@@ -27,8 +34,13 @@ def dataloader(batch_size: int = 25) -> DataLoader:
     return to_dataloader(x, y, batch_size=batch_size, infinite=True)
 
 
-data = dataloader()
-model = QNN(QuantumCircuit(2, hea(2, 1)), observable=Z(0))
+data = dataloader(hyperparams["batch_size"])
+model = QNN(
+    QuantumCircuit(
+        hyperparams["n_qubits"], hea(hyperparams["n_qubits"], hyperparams["ansatz_depth"])
+    ),
+    observable=hyperparams["observable"](0),
+)
 cnt = count()
 criterion = torch.nn.MSELoss()
 optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
@@ -48,6 +60,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     checkpoint_every=1,
     write_every=1,
     tracking_tool=ExperimentTrackingTool.MLFLOW,
+    hyperparams=hyperparams,
 )
 train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
 
diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 39f161dc2..7e10c15a3 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -53,6 +53,9 @@ class TrainConfig:
     verbose: bool = True
     """Whether or not to print out metrics values during training."""
     tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+    """The tracking tool of choice."""
+    hyperparams: Optional[dict] = None
+    """Hyperparameters to track."""
 
     # mlflow_callbacks: list[Callable] = [write_mlflow_figure(), write_x()]
 
diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 11aba23db..7a2f0e778 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -23,14 +23,17 @@ 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 write_mlflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
-    # TODO for giorgio
     writer.log_metrics({"loss": float(loss)}, step=iteration)  # type: ignore
-    writer.log_metrics(metrics, step=iteration)
+    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)
 
 
 TRACKER_MAPPING = {
@@ -38,8 +41,19 @@ def write_mlflow(writer: Any, loss: float | None, metrics: dict, iteration: int)
     ExperimentTrackingTool.MLFLOW: write_mlflow,
 }
 
+LOGGER_MAPPING = {
+    ExperimentTrackingTool.TENSORBOARD: log_hyperparams_tensorboard,
+    ExperimentTrackingTool.MLFLOW: log_hyperparams_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)
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index fdca919ac..fe75e4a63 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -4,7 +4,13 @@
 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,7 +22,7 @@
 from qadence.ml_tools.config import MLFlowConfig, TrainConfig
 from qadence.ml_tools.data import DictDataLoader
 from qadence.ml_tools.optimize_step import optimize_step
-from qadence.ml_tools.printing import print_metrics, write_tracker
+from qadence.ml_tools.printing import log_tracker, print_metrics, write_tracker
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
 from qadence.types import ExperimentTrackingTool
 
@@ -31,7 +37,6 @@ def train(
     loss_fn: Callable,
     device: torch_device = None,
     optimize_step: Callable = optimize_step,
-    write_tracker: Callable = write_tracker,
     dtype: torch_dtype = None,
 ) -> tuple[Module, Optimizer]:
     """Runs the training loop with gradient-based optimizer.
@@ -134,6 +139,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
         # writer.mlflow.pytorch.autolog(
         #     log_every_n_step=config.write_every, log_models=False, log_datasets=False
         # )
+
     ## Training
     progress = Progress(
         TextColumn("[progress.description]{task.description}"),
@@ -194,6 +200,10 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                 logger.info("Terminating training gracefully after the current iteration.")
                 break
 
+    # writing hyperparameters
+    if config.hyperparams:
+        log_tracker((writer, config.hyperparams, metrics), tracking_tool=config.tracking_tool)
+
     # Final writing and checkpointing
     if config.folder:
         write_checkpoint(config.folder, model, optimizer, iteration)
@@ -201,6 +211,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer.close()
     elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:
+        # track some final stuff e.g. images
         writer.end_run()
 
     return model, optimizer

From d12bc249533887d04fc25dd457d267981c6a8232 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Thu, 20 Jun 2024 17:28:25 +0200
Subject: [PATCH 10/51] Add figures logging with callback

---
 docs/tutorials/qml/mlflow_demonstration.py | 29 ++++++++++++++++--
 qadence/ml_tools/config.py                 | 14 +++++++++
 qadence/ml_tools/printing.py               | 34 ++++++++++++++++++++--
 qadence/ml_tools/train_grad.py             | 12 +++++++-
 4 files changed, 84 insertions(+), 5 deletions(-)

diff --git a/docs/tutorials/qml/mlflow_demonstration.py b/docs/tutorials/qml/mlflow_demonstration.py
index ab7cc2459..7922e0861 100644
--- a/docs/tutorials/qml/mlflow_demonstration.py
+++ b/docs/tutorials/qml/mlflow_demonstration.py
@@ -1,12 +1,17 @@
 from __future__ import annotations
 
 import os
+import random
 from itertools import count
 
+import numpy as np
 import torch
+from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
 from torch.utils.data import DataLoader
 
 from qadence import QuantumCircuit, Z, hea
+from qadence.constructors import feature_map, hamiltonian_factory
 from qadence.ml_tools import TrainConfig, train_with_grad
 from qadence.ml_tools.data import to_dataloader
 from qadence.ml_tools.utils import rand_featureparameters
@@ -18,12 +23,17 @@
 os.environ["MLFLOW_RUN_NAME"] = "test_0"
 
 hyperparams = {
+    "seed": 42,
     "batch_size": 10,
     "n_qubits": 2,
     "ansatz_depth": 1,
     "observable": Z,
 }
 
+np.random.seed(hyperparams["seed"])
+torch.manual_seed(hyperparams["seed"])
+random.seed(hyperparams["seed"])
+
 
 # in case you want to track remotely
 # os.environ['MLFLOW_TRACKING_USERNAME'] =
@@ -34,12 +44,16 @@ def dataloader(batch_size: int = 25) -> DataLoader:
     return to_dataloader(x, y, batch_size=batch_size, infinite=True)
 
 
+obs = hamiltonian_factory(register=hyperparams["n_qubits"], detuning=hyperparams["observable"])
+
 data = dataloader(hyperparams["batch_size"])
+fm = feature_map(hyperparams["n_qubits"], param="x")
 model = QNN(
     QuantumCircuit(
-        hyperparams["n_qubits"], hea(hyperparams["n_qubits"], hyperparams["ansatz_depth"])
+        hyperparams["n_qubits"], fm, hea(hyperparams["n_qubits"], hyperparams["ansatz_depth"])
     ),
-    observable=hyperparams["observable"](0),
+    observable=obs,
+    inputs=["x"],
 )
 cnt = count()
 criterion = torch.nn.MSELoss()
@@ -54,13 +68,24 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     return loss, {}
 
 
+def plot_fn(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
+    descr = f"ufa_prediction_epoch_{iteration}.png"
+    fig, ax = plt.subplots()
+    x = torch.linspace(0, 1, 100).reshape(-1, 1)
+    out = model.expectation(x)
+    ax.plot(x.detach().numpy(), out.detach().numpy())
+    return descr, fig
+
+
 config = TrainConfig(
     folder="mlflow_demonstration",
     max_iter=10,
     checkpoint_every=1,
+    plot_every=2,
     write_every=1,
     tracking_tool=ExperimentTrackingTool.MLFLOW,
     hyperparams=hyperparams,
+    plotting_functions=(plot_fn,),
 )
 train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
 
diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 7e10c15a3..24fc01860 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -8,6 +8,9 @@
 from typing import Callable, Optional
 from uuid import uuid4
 
+from matplotlib.figure import Figure
+from torch.nn import Module
+
 from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
@@ -35,6 +38,11 @@ class TrainConfig:
     """Write tensorboard logs."""
     checkpoint_every: int = 5000
     """Write model/optimizer checkpoint."""
+    plot_every: Optional[int] = None
+    """Write figures.
+
+    NOTE: currently only works with mlflow.
+    """
     folder: Optional[Path] = None
     """Checkpoint/tensorboard logs folder."""
     create_subfolder_per_run: bool = False
@@ -56,6 +64,8 @@ class TrainConfig:
     """The tracking tool of choice."""
     hyperparams: Optional[dict] = None
     """Hyperparameters to track."""
+    plotting_functions: Optional[tuple[Callable[[Module, int], tuple[str, Figure]]]] = None
+    """Functions for in-train plotting."""
 
     # mlflow_callbacks: list[Callable] = [write_mlflow_figure(), write_x()]
 
@@ -72,6 +82,10 @@ def __post_init__(self) -> None:
             self.trainstop_criterion = lambda x: x <= self.max_iter
         if self.validation_criterion is None:
             self.validation_criterion = lambda x: False
+        if self.plot_every and self.tracking_tool != ExperimentTrackingTool.MLFLOW:
+            raise NotImplementedError("In-training plots are only available with mlflow tracking.")
+        if self.plot_every and self.plotting_functions is None:
+            logger.warning("Plots tracking is required, but no plotting functions are provided.")
 
 
 @dataclass
diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 7a2f0e778..09ee39d86 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -1,7 +1,9 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Callable
 
+from matplotlib.figure import Figure
+from torch.nn import Module
 from torch.utils.tensorboard import SummaryWriter
 
 from qadence.types import ExperimentTrackingTool
@@ -27,13 +29,30 @@ def log_hyperparams_tensorboard(writer: SummaryWriter, hyperparams: dict, metric
     writer.add_hparams(hyperparams, metrics)
 
 
+def plot_tensorboard(
+    writer: SummaryWriter, iteration: int, plotting_functions: tuple[Callable]
+) -> None:
+    raise NotImplementedError("Plot logging with tensorboard is not implemented")
+
+
 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)
+    writer.log_params(hyperparams)  # type: ignore
+
+
+def plot_mlflow(
+    writer: SummaryWriter,
+    model: Module,
+    iteration: int,
+    plotting_functions: tuple[Callable[[Module, int], tuple[str, Figure]]],
+) -> None:
+    for pf in plotting_functions:
+        descr, fig = pf(model, iteration)
+        writer.log_figure(fig, descr)
 
 
 TRACKER_MAPPING = {
@@ -46,6 +65,11 @@ def log_hyperparams_mlflow(writer: Any, hyperparams: dict, metrics: dict) -> Non
     ExperimentTrackingTool.MLFLOW: log_hyperparams_mlflow,
 }
 
+PLOTTER_MAPPING = {
+    ExperimentTrackingTool.TENSORBOARD: plot_tensorboard,
+    ExperimentTrackingTool.MLFLOW: plot_mlflow,
+}
+
 
 def write_tracker(
     args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
@@ -57,3 +81,9 @@ 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)  # type: ignore
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index fe75e4a63..bdc6df7e1 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -22,7 +22,12 @@
 from qadence.ml_tools.config import MLFlowConfig, TrainConfig
 from qadence.ml_tools.data import DictDataLoader
 from qadence.ml_tools.optimize_step import optimize_step
-from qadence.ml_tools.printing import log_tracker, print_metrics, write_tracker
+from qadence.ml_tools.printing import (
+    log_tracker,
+    plot_tracker,
+    print_metrics,
+    write_tracker,
+)
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
 from qadence.types import ExperimentTrackingTool
 
@@ -192,6 +197,11 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                 if iteration % config.write_every == 0:
                     write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
 
+                if iteration % config.plot_every == 0:
+                    plot_tracker(
+                        (writer, model, iteration, config.plotting_functions), config.tracking_tool
+                    )
+
                 if config.folder:
                     if iteration % config.checkpoint_every == 0:
                         write_checkpoint(config.folder, model, optimizer, iteration)

From 88195ab9ccdb84b1e451681771aa638b9b83bc91 Mon Sep 17 00:00:00 2001
From: smit <smit.chaudhary@pasqal.com>
Date: Thu, 4 Jul 2024 10:48:45 +0200
Subject: [PATCH 11/51] Add details to MLFlowConfig class

---
 qadence/ml_tools/config.py | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 24fc01860..36877c91b 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -38,12 +38,12 @@ class TrainConfig:
     """Write tensorboard logs."""
     checkpoint_every: int = 5000
     """Write model/optimizer checkpoint."""
-    plot_every: Optional[int] = None
+    plot_every: int | None = None
     """Write figures.
 
     NOTE: currently only works with mlflow.
     """
-    folder: Optional[Path] = None
+    folder: Path | None = None
     """Checkpoint/tensorboard logs folder."""
     create_subfolder_per_run: bool = False
     """Checkpoint/tensorboard logs stored in subfolder with name `<timestamp>_<PID>`.
@@ -52,7 +52,7 @@ class TrainConfig:
     """
     checkpoint_best_only: bool = False
     """Write model/optimizer checkpoint only if a metric has improved."""
-    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
     """A boolean function which evaluates a given training stopping metric is satisfied."""
@@ -62,9 +62,9 @@ class TrainConfig:
     """Whether or not to print out metrics values during training."""
     tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
     """The tracking tool of choice."""
-    hyperparams: Optional[dict] = None
+    hyperparams: dict | None = None
     """Hyperparameters to track."""
-    plotting_functions: Optional[tuple[Callable[[Module, int], tuple[str, Figure]]]] = None
+    plotting_functions: tuple[Callable[[Module, int], tuple[str, Figure]]] | None = None
     """Functions for in-train plotting."""
 
     # mlflow_callbacks: list[Callable] = [write_mlflow_figure(), write_x()]
@@ -84,13 +84,15 @@ def __post_init__(self) -> None:
             self.validation_criterion = lambda x: False
         if self.plot_every and self.tracking_tool != ExperimentTrackingTool.MLFLOW:
             raise NotImplementedError("In-training plots are only available with mlflow tracking.")
-        if self.plot_every and self.plotting_functions is None:
+        if self.plot_every is not None and self.plotting_functions is None:
             logger.warning("Plots tracking is required, but no plotting functions are provided.")
 
 
 @dataclass
 class MLFlowConfig:
     """
+    Configuration for mlflow tracking.
+
     Example:
 
         export MLFLOW_TRACKING_URI=tracking_uri
@@ -99,10 +101,26 @@ class MLFlowConfig:
     """
 
     MLFLOW_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).
+    """
+
     MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
+    """The username for the mlflow tracking server."""
+
     MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
+    """The password for the mlflow tracking server."""
+
     EXPERIMENT: 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.
+    """
+
     RUN_NAME: str = os.getenv("MLFLOW_RUN_NAME", "test_0")
+    """The name of the run."""
 
     def __post_init__(self) -> None:
         import mlflow

From 72f109872e3ac926e6abb82c594c60b288955947 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Tue, 9 Jul 2024 18:35:45 +0200
Subject: [PATCH 12/51] Sync tracking stuff in train_no_grad

---
 qadence/ml_tools/train_no_grad.py | 44 +++++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 8 deletions(-)

diff --git a/qadence/ml_tools/train_no_grad.py b/qadence/ml_tools/train_no_grad.py
index d50250d2f..03fdd3093 100644
--- a/qadence/ml_tools/train_no_grad.py
+++ b/qadence/ml_tools/train_no_grad.py
@@ -1,22 +1,34 @@
 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
 from torch.utils.tensorboard import SummaryWriter
 
-from qadence.ml_tools.config import TrainConfig
+from qadence.ml_tools.config import MLFlowConfig, 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 (
+    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__)
 
@@ -63,8 +75,16 @@ def _update_parameters(
     # 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:
+        MLFlowConfig()  # Set up credentials for mlflow tracking
+        writer = importlib.import_module("mlflow")
+
+        # writer.mlflow.pytorch.autolog(
+        #     log_every_n_step=config.write_every, log_models=False, log_datasets=False
+        # )
 
     # set optimizer configuration and initial parameters
     optimizer.budget = config.max_iter
@@ -100,7 +120,12 @@ def _update_parameters(
                 print_metrics(loss, metrics, iteration)
 
             if iteration % config.write_every == 0:
-                write_tensorboard(writer, loss, metrics, iteration)
+                write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+
+            if iteration % config.plot_every == 0:
+                plot_tracker(
+                    (writer, model, iteration, config.plotting_functions), config.tracking_tool
+                )
 
             if config.folder:
                 if iteration % config.checkpoint_every == 0:
@@ -112,7 +137,10 @@ def _update_parameters(
     ## Final writing and stuff
     if config.folder:
         write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tensorboard(writer, loss, metrics, iteration)
-    writer.close()
+    write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer.close()
+    elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:
+        writer.end_run()
 
     return model, optimizer

From 936021e4b593ad975eeb6f75a07ab44ff319ebe7 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 10 Jul 2024 09:52:14 +0200
Subject: [PATCH 13/51] Fix current ml tracking tests

---
 qadence/ml_tools/config.py   | 17 +++++++----------
 qadence/ml_tools/printing.py | 13 ++++++++-----
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 36877c91b..89053cad8 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -2,15 +2,12 @@
 
 import datetime
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from logging import getLogger
 from pathlib import Path
 from typing import Callable, Optional
 from uuid import uuid4
 
-from matplotlib.figure import Figure
-from torch.nn import Module
-
 from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
@@ -38,7 +35,7 @@ class TrainConfig:
     """Write tensorboard logs."""
     checkpoint_every: int = 5000
     """Write model/optimizer checkpoint."""
-    plot_every: int | None = None
+    plot_every: int = 5000
     """Write figures.
 
     NOTE: currently only works with mlflow.
@@ -64,7 +61,7 @@ class TrainConfig:
     """The tracking tool of choice."""
     hyperparams: dict | None = None
     """Hyperparameters to track."""
-    plotting_functions: tuple[Callable[[Module, int], tuple[str, Figure]]] | None = None
+    plotting_functions: tuple[Callable] = field(default_factory=tuple)  # type: ignore
     """Functions for in-train plotting."""
 
     # mlflow_callbacks: list[Callable] = [write_mlflow_figure(), write_x()]
@@ -82,10 +79,10 @@ def __post_init__(self) -> None:
             self.trainstop_criterion = lambda x: x <= self.max_iter
         if self.validation_criterion is None:
             self.validation_criterion = lambda x: False
-        if self.plot_every and self.tracking_tool != ExperimentTrackingTool.MLFLOW:
-            raise NotImplementedError("In-training plots are only available with mlflow tracking.")
-        if self.plot_every is not None and self.plotting_functions is None:
-            logger.warning("Plots tracking is required, but no plotting functions are provided.")
+        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.")
 
 
 @dataclass
diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 09ee39d86..ed833307e 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -2,7 +2,6 @@
 
 from typing import Any, Callable
 
-from matplotlib.figure import Figure
 from torch.nn import Module
 from torch.utils.tensorboard import SummaryWriter
 
@@ -30,9 +29,13 @@ def log_hyperparams_tensorboard(writer: SummaryWriter, hyperparams: dict, metric
 
 
 def plot_tensorboard(
-    writer: SummaryWriter, iteration: int, plotting_functions: tuple[Callable]
+    writer: SummaryWriter,
+    model: Module,
+    iteration: int,
+    plotting_functions: tuple[Callable],
 ) -> None:
-    raise NotImplementedError("Plot logging with tensorboard is not implemented")
+    # TODO: implement me
+    pass
 
 
 def write_mlflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
@@ -45,10 +48,10 @@ def log_hyperparams_mlflow(writer: Any, hyperparams: dict, metrics: dict) -> Non
 
 
 def plot_mlflow(
-    writer: SummaryWriter,
+    writer: Any,
     model: Module,
     iteration: int,
-    plotting_functions: tuple[Callable[[Module, int], tuple[str, Figure]]],
+    plotting_functions: tuple[Callable],
 ) -> None:
     for pf in plotting_functions:
         descr, fig = pf(model, iteration)

From 8d630ccdd8d7aba421f715e5e734cfbb68bc4e2d Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 10 Jul 2024 11:49:53 +0200
Subject: [PATCH 14/51] Add hparams filtering for TB

---
 qadence/ml_tools/config.py | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 89053cad8..b5a215b44 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -8,6 +8,8 @@
 from typing import Callable, Optional
 from uuid import uuid4
 
+from torch import Tensor
+
 from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
@@ -59,12 +61,28 @@ class TrainConfig:
     """Whether or not to print out metrics values during training."""
     tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
     """The tracking tool of choice."""
-    hyperparams: dict | None = None
+    hyperparams: dict = field(default_factory=dict)
     """Hyperparameters to track."""
     plotting_functions: tuple[Callable] = field(default_factory=tuple)  # type: ignore
     """Functions for in-train plotting."""
 
-    # mlflow_callbacks: list[Callable] = [write_mlflow_figure(), write_x()]
+    # 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:
@@ -79,6 +97,8 @@ def __post_init__(self) -> None:
             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.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:

From eaab95a666ee406d9102d7f255897fc53dae5ef1 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 10 Jul 2024 12:03:30 +0200
Subject: [PATCH 15/51] Test all current on mlflow too

---
 tests/ml_tools/test_checkpointing.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/tests/ml_tools/test_checkpointing.py b/tests/ml_tools/test_checkpointing.py
index 61e0a96ad..5ba78560c 100644
--- a/tests/ml_tools/test_checkpointing.py
+++ b/tests/ml_tools/test_checkpointing.py
@@ -79,7 +79,12 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     assert torch.allclose(loaded_model.expectation({}), model.expectation({}))
 
 
-def test_check_ckpts_exist(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
+@pytest.mark.parametrize(
+    "tool", [ExperimentTrackingTool.TENSORBOARD, ExperimentTrackingTool.MLFLOW]
+)
+def test_check_ckpts_exist(
+    tool: ExperimentTrackingTool, BasicQuantumModel: QuantumModel, tmp_path: Path
+) -> None:
     data = dataloader()
     model = BasicQuantumModel
     cnt = count()
@@ -92,7 +97,9 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         loss = criterion(out, torch.rand(1))
         return loss, {}
 
-    config = TrainConfig(folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1)
+    config = TrainConfig(
+        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=tool
+    )
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
     ckpts = [tmp_path / Path(f"model_QuantumModel_ckpt_00{i}_device_cpu.pt") for i in range(1, 9)]
     assert all(os.path.isfile(ckpt) for ckpt in ckpts)
@@ -189,8 +196,11 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     assert torch.allclose(loaded_model.expectation(inputs), model.expectation(inputs))
 
 
+@pytest.mark.parametrize(
+    "tool", [ExperimentTrackingTool.TENSORBOARD, ExperimentTrackingTool.MLFLOW]
+)
 def test_check_transformedmodule_ckpts_exist(
-    BasicTransformedModule: TransformedModule, tmp_path: Path
+    tool: ExperimentTrackingTool, BasicTransformedModule: TransformedModule, tmp_path: Path
 ) -> None:
     data = dataloader()
     model = BasicTransformedModule
@@ -205,7 +215,9 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         loss = criterion(out, torch.rand(1))
         return loss, {}
 
-    config = TrainConfig(folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1)
+    config = TrainConfig(
+        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=tool
+    )
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
     ckpts = [
         tmp_path / Path(f"model_TransformedModule_ckpt_00{i}_device_cpu.pt") for i in range(1, 9)

From d9d8b48ab25794bb660fcfadb93e32f3d6d95fe2 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 10 Jul 2024 16:37:08 +0200
Subject: [PATCH 16/51] Correct experiment creation/setup

---
 qadence/ml_tools/config.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index b5a215b44..909042cfd 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -147,5 +147,10 @@ def __post_init__(self) -> None:
                 f"Intialized mlflow remote logging for user {self.MLFLOW_TRACKING_USERNAME}."
             )
         mlflow.set_tracking_uri(self.MLFLOW_TRACKING_URI)
+        # activate existing or create experiment
+        exp_filter_string = f"name = '{self.EXPERIMENT}'"
+        if not mlflow.search_experiments(filter_string=exp_filter_string):
+            mlflow.create_experiment(name=self.EXPERIMENT)
+
         mlflow.set_experiment(self.EXPERIMENT)
         mlflow.start_run(run_name=self.RUN_NAME, nested=False)

From 0211ef82e16473d8c7ccfde999ce60a973098b65 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 10 Jul 2024 16:38:27 +0200
Subject: [PATCH 17/51] Remove os.environ (vars do not persist bc of imorts)

---
 docs/tutorials/qml/mlflow_demonstration.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/docs/tutorials/qml/mlflow_demonstration.py b/docs/tutorials/qml/mlflow_demonstration.py
index 7922e0861..f27002712 100644
--- a/docs/tutorials/qml/mlflow_demonstration.py
+++ b/docs/tutorials/qml/mlflow_demonstration.py
@@ -18,10 +18,6 @@
 from qadence.models import QNN, QuantumModel
 from qadence.types import ExperimentTrackingTool
 
-os.environ["MLFLOW_TRACKING_URI"] = "sqlite:///mlflow.db"
-os.environ["MLFLOW_EXPERIMENT"] = "mlflow_demonstration"
-os.environ["MLFLOW_RUN_NAME"] = "test_0"
-
 hyperparams = {
     "seed": 42,
     "batch_size": 10,
@@ -87,7 +83,7 @@ def plot_fn(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
     hyperparams=hyperparams,
     plotting_functions=(plot_fn,),
 )
+
 train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
 
 os.system("mlflow ui --port 5000")
-os.system("mlflow ui --backend-store-uri sqlite:///mlflow.db")

From d126505305efb2a58b2edf4df70ff3027c808d69 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Thu, 11 Jul 2024 06:53:30 +0200
Subject: [PATCH 18/51] Ignore artifacts

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index a18f2e0e4..b066019b6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -78,3 +78,4 @@ events.out.tfevents.*
 
 # mlflow
 mlruns/
+mlartifacts/

From 039ac8ff0ec37256d267e09aac0ac443f90d1f14 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Fri, 12 Jul 2024 08:52:52 +0200
Subject: [PATCH 19/51] Move MLFlowConfig to TrainConfig

---
 qadence/ml_tools/config.py        | 21 ++++++++++++++++-----
 qadence/ml_tools/train_grad.py    |  3 +--
 qadence/ml_tools/train_no_grad.py |  9 ++-------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 909042cfd..b6a875ce0 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -99,11 +99,22 @@ def __post_init__(self) -> 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."
+            )
+
 
 @dataclass
 class MLFlowConfig:
@@ -130,7 +141,7 @@ class MLFlowConfig:
     MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
     """The password for the mlflow tracking server."""
 
-    EXPERIMENT: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
+    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.
@@ -148,9 +159,9 @@ def __post_init__(self) -> None:
             )
         mlflow.set_tracking_uri(self.MLFLOW_TRACKING_URI)
         # activate existing or create experiment
-        exp_filter_string = f"name = '{self.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)
+            mlflow.create_experiment(name=self.EXPERIMENT_NAME)
 
-        mlflow.set_experiment(self.EXPERIMENT)
-        mlflow.start_run(run_name=self.RUN_NAME, nested=False)
+        self.experiment = mlflow.set_experiment(self.EXPERIMENT_NAME)
+        self.run = mlflow.start_run(run_name=self.RUN_NAME, nested=False)
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index bdc6df7e1..71963eab7 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -19,7 +19,7 @@
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
-from qadence.ml_tools.config import MLFlowConfig, TrainConfig
+from qadence.ml_tools.config import TrainConfig
 from qadence.ml_tools.data import DictDataLoader
 from qadence.ml_tools.optimize_step import optimize_step
 from qadence.ml_tools.printing import (
@@ -138,7 +138,6 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer = SummaryWriter(config.folder, purge_step=init_iter)
     else:
-        MLFlowConfig()  # Set up credentials for mlflow tracking
         writer = importlib.import_module("mlflow")
 
         # writer.mlflow.pytorch.autolog(
diff --git a/qadence/ml_tools/train_no_grad.py b/qadence/ml_tools/train_no_grad.py
index 03fdd3093..1b89dfbf6 100644
--- a/qadence/ml_tools/train_no_grad.py
+++ b/qadence/ml_tools/train_no_grad.py
@@ -18,14 +18,10 @@
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
-from qadence.ml_tools.config import MLFlowConfig, TrainConfig
+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 (
-    plot_tracker,
-    print_metrics,
-    write_tracker,
-)
+from qadence.ml_tools.printing import 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
@@ -79,7 +75,6 @@ def _update_parameters(
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer = SummaryWriter(config.folder, purge_step=init_iter)
     else:
-        MLFlowConfig()  # Set up credentials for mlflow tracking
         writer = importlib.import_module("mlflow")
 
         # writer.mlflow.pytorch.autolog(

From 1e560a0d7844a39970544dde18d9d81f7ebaf916 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Fri, 12 Jul 2024 08:59:38 +0200
Subject: [PATCH 20/51] De-dataclass MLConfig

Reason: dataclass could not read env vars set by os.environ, making it hard to test
---
 qadence/ml_tools/config.py | 35 +++++++++++++++++------------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index b6a875ce0..8208930a7 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -116,7 +116,6 @@ def mlflow_config(self) -> MLFlowConfig:
             )
 
 
-@dataclass
 class MLFlowConfig:
     """
     Configuration for mlflow tracking.
@@ -128,30 +127,30 @@ class MLFlowConfig:
         export MLFLOW_TRACKING_PASSWORD=password
     """
 
-    MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
-    """The URI of the mlflow tracking server.
+    def __init__(self) -> None:
+        import mlflow
 
-    An empty string, or a local file path, prefixed with file:/.
-    Data is stored locally at the provided file (or ./mlruns if empty).
-    """
+        self.MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
+        """The URI of the mlflow tracking server.
 
-    MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
-    """The username for 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).
+        """
 
-    MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
-    """The password for the mlflow tracking server."""
+        self.MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
+        """The username for the mlflow tracking server."""
 
-    EXPERIMENT_NAME: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
-    """The name of the experiment.
+        self.MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
+        """The password for the mlflow tracking server."""
 
-    If None or empty, a new experiment is created with a random UUID.
-    """
+        self.EXPERIMENT_NAME: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
+        """The name of the experiment.
 
-    RUN_NAME: str = os.getenv("MLFLOW_RUN_NAME", "test_0")
-    """The name of the run."""
+        If None or empty, a new experiment is created with a random UUID.
+        """
 
-    def __post_init__(self) -> None:
-        import mlflow
+        self.RUN_NAME: str = os.getenv("MLFLOW_RUN_NAME", "test_0")
+        """The name of the run."""
 
         if self.MLFLOW_TRACKING_USERNAME != "":
             logger.info(

From 0ba4c2c99f515dc11aaf7d0eae578a5d2258c210 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Fri, 12 Jul 2024 09:37:19 +0200
Subject: [PATCH 21/51] Lowercase variables in MLFlowConfig

---
 qadence/ml_tools/config.py | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 8208930a7..4532ed480 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -130,37 +130,36 @@ class MLFlowConfig:
     def __init__(self) -> None:
         import mlflow
 
-        self.MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
+        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.MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
+        self.tracking_username: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
         """The username for the mlflow tracking server."""
 
-        self.MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
+        self.tracking_pwd: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
         """The password for the mlflow tracking server."""
 
-        self.EXPERIMENT_NAME: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
+        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", "test_0")
+        self.run_name: str = os.getenv("MLFLOW_RUN_NAME", "test_0")
         """The name of the run."""
 
-        if self.MLFLOW_TRACKING_USERNAME != "":
-            logger.info(
-                f"Intialized mlflow remote logging for user {self.MLFLOW_TRACKING_USERNAME}."
-            )
-        mlflow.set_tracking_uri(self.MLFLOW_TRACKING_URI)
+        if self.tracking_username != "":
+            logger.info(f"Intialized mlflow remote logging for user {self.tracking_username}.")
+
+        mlflow.set_tracking_uri(self.tracking_uri)
         # activate existing or create experiment
-        exp_filter_string = f"name = '{self.EXPERIMENT_NAME}'"
+        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)
+            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)
+        self.experiment = mlflow.set_experiment(self.experiment_name)
+        self.run = mlflow.start_run(run_name=self.run_name, nested=False)

From d412965619cd235aa052cc48947dad827678e5dd Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Fri, 12 Jul 2024 11:32:12 +0200
Subject: [PATCH 22/51] Adjust tests and introduce test_logging

---
 tests/ml_tools/test_checkpointing.py | 31 +++------------
 tests/ml_tools/test_logging.py       | 58 ++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 25 deletions(-)
 create mode 100644 tests/ml_tools/test_logging.py

diff --git a/tests/ml_tools/test_checkpointing.py b/tests/ml_tools/test_checkpointing.py
index 5ba78560c..627dea7de 100644
--- a/tests/ml_tools/test_checkpointing.py
+++ b/tests/ml_tools/test_checkpointing.py
@@ -4,7 +4,6 @@
 from itertools import count
 from pathlib import Path
 
-import pytest
 import torch
 from torch.utils.data import DataLoader
 
@@ -19,7 +18,6 @@
 from qadence.ml_tools.parameters import get_parameters, set_parameters
 from qadence.ml_tools.utils import rand_featureparameters
 from qadence.models import QNN, QuantumModel
-from qadence.types import ExperimentTrackingTool
 
 
 def dataloader(batch_size: int = 25) -> DataLoader:
@@ -79,12 +77,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     assert torch.allclose(loaded_model.expectation({}), model.expectation({}))
 
 
-@pytest.mark.parametrize(
-    "tool", [ExperimentTrackingTool.TENSORBOARD, ExperimentTrackingTool.MLFLOW]
-)
-def test_check_ckpts_exist(
-    tool: ExperimentTrackingTool, BasicQuantumModel: QuantumModel, tmp_path: Path
-) -> None:
+def test_check_ckpts_exist(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
     data = dataloader()
     model = BasicQuantumModel
     cnt = count()
@@ -97,9 +90,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         loss = criterion(out, torch.rand(1))
         return loss, {}
 
-    config = TrainConfig(
-        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=tool
-    )
+    config = TrainConfig(folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1)
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
     ckpts = [tmp_path / Path(f"model_QuantumModel_ckpt_00{i}_device_cpu.pt") for i in range(1, 9)]
     assert all(os.path.isfile(ckpt) for ckpt in ckpts)
@@ -138,10 +129,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     assert torch.allclose(loaded_model.expectation(inputs), model.expectation(inputs))
 
 
-@pytest.mark.parametrize(
-    "tool", [ExperimentTrackingTool.TENSORBOARD, ExperimentTrackingTool.MLFLOW]
-)
-def test_check_QNN_ckpts_exist(tool: ExperimentTrackingTool, BasicQNN: QNN, tmp_path: Path) -> None:
+def test_check_QNN_ckpts_exist(BasicQNN: QNN, tmp_path: Path) -> None:
     data = dataloader()
     model = BasicQNN
     cnt = count()
@@ -155,9 +143,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         loss = criterion(out, torch.rand(1))
         return loss, {}
 
-    config = TrainConfig(
-        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=tool
-    )
+    config = TrainConfig(folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1)
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
     ckpts = [tmp_path / Path(f"model_QNN_ckpt_00{i}_device_cpu.pt") for i in range(1, 9)]
     assert all(os.path.isfile(ckpt) for ckpt in ckpts)
@@ -196,11 +182,8 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     assert torch.allclose(loaded_model.expectation(inputs), model.expectation(inputs))
 
 
-@pytest.mark.parametrize(
-    "tool", [ExperimentTrackingTool.TENSORBOARD, ExperimentTrackingTool.MLFLOW]
-)
 def test_check_transformedmodule_ckpts_exist(
-    tool: ExperimentTrackingTool, BasicTransformedModule: TransformedModule, tmp_path: Path
+    BasicTransformedModule: TransformedModule, tmp_path: Path
 ) -> None:
     data = dataloader()
     model = BasicTransformedModule
@@ -215,9 +198,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         loss = criterion(out, torch.rand(1))
         return loss, {}
 
-    config = TrainConfig(
-        folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=tool
-    )
+    config = TrainConfig(folder=tmp_path, max_iter=10, checkpoint_every=1, write_every=1)
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
     ckpts = [
         tmp_path / Path(f"model_TransformedModule_ckpt_00{i}_device_cpu.pt") for i in range(1, 9)
diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
new file mode 100644
index 000000000..f750b059a
--- /dev/null
+++ b/tests/ml_tools/test_logging.py
@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import os
+from itertools import count
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from qadence.ml_tools import TrainConfig, train_with_grad
+from qadence.ml_tools.data import to_dataloader
+from qadence.models import QuantumModel
+from qadence.types import ExperimentTrackingTool
+
+
+def dataloader(batch_size: int = 25) -> DataLoader:
+    x = torch.linspace(0, 1, batch_size).reshape(-1, 1)
+    y = torch.cos(x)
+    return to_dataloader(x, y, batch_size=batch_size, infinite=True)
+
+
+def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
+    data = dataloader()
+    model = BasicQuantumModel
+    cnt = count()
+    criterion = torch.nn.MSELoss()
+    hyperparams = {"max_iter": int(10), "lr": 0.1}
+    optimizer = torch.optim.Adam(model.parameters(), lr=hyperparams["lr"])
+
+    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+        next(cnt)
+        out = model.expectation({})
+        loss = criterion(out, torch.rand(1))
+        return loss, {}
+
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=hyperparams["max_iter"],  # type: ignore
+        checkpoint_every=1,
+        write_every=1,
+        hyperparams=hyperparams,
+        tracking_tool=ExperimentTrackingTool.MLFLOW,
+    )
+
+    train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+    mlflow_config = config.mlflow_config
+    experiment_id = mlflow_config.run.info.experiment_id
+    run_id = mlflow_config.run.info.run_id
+
+    hyperparams_files = [
+        Path(f"mlruns/{experiment_id}/{run_id}/params/{key}") for key in hyperparams.keys()
+    ]
+    assert all([os.path.isfile(hf) for hf in hyperparams_files])
+
+
+def test_plotting() -> None:
+    pass

From 714939a8db13fb94b54173795f023b397c3c6373 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Mon, 15 Jul 2024 11:05:35 +0200
Subject: [PATCH 23/51] Add mlflow final model logging

---
 docs/tutorials/qml/mlflow_demonstration.py |  8 ++--
 qadence/ml_tools/config.py                 |  2 +
 qadence/ml_tools/printing.py               | 51 +++++++++++++++++++++-
 qadence/ml_tools/train_grad.py             |  5 ++-
 qadence/ml_tools/train_no_grad.py          | 19 +++++++-
 5 files changed, 78 insertions(+), 7 deletions(-)

diff --git a/docs/tutorials/qml/mlflow_demonstration.py b/docs/tutorials/qml/mlflow_demonstration.py
index f27002712..f34a94d72 100644
--- a/docs/tutorials/qml/mlflow_demonstration.py
+++ b/docs/tutorials/qml/mlflow_demonstration.py
@@ -1,11 +1,12 @@
 from __future__ import annotations
 
-import os
 import random
 from itertools import count
 
 import numpy as np
 import torch
+
+# from dotenv import load_dotenv
 from matplotlib import pyplot as plt
 from matplotlib.figure import Figure
 from torch.utils.data import DataLoader
@@ -30,6 +31,8 @@
 torch.manual_seed(hyperparams["seed"])
 random.seed(hyperparams["seed"])
 
+# load_dotenv()
+
 
 # in case you want to track remotely
 # os.environ['MLFLOW_TRACKING_USERNAME'] =
@@ -79,11 +82,10 @@ def plot_fn(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
     checkpoint_every=1,
     plot_every=2,
     write_every=1,
+    log_model=True,
     tracking_tool=ExperimentTrackingTool.MLFLOW,
     hyperparams=hyperparams,
     plotting_functions=(plot_fn,),
 )
 
 train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
-
-os.system("mlflow ui --port 5000")
diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 4532ed480..13c48d02e 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -42,6 +42,8 @@ class TrainConfig:
 
     NOTE: currently only works with mlflow.
     """
+    log_model: bool = True
+    """Logs a serialised version of the model."""
     folder: Path | None = None
     """Checkpoint/tensorboard logs folder."""
     create_subfolder_per_run: bool = False
diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index ed833307e..be0d92a68 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -1,10 +1,13 @@
 from __future__ import annotations
 
-from typing import Any, Callable
+from typing import Any, Callable, Sequence, Union
 
+from mlflow.models import infer_signature
 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
 
 
@@ -38,6 +41,15 @@ def plot_tensorboard(
     pass
 
 
+def log_model_tensorboard(
+    writer: SummaryWriter,
+    model: Module,
+    dataloader: Union[None, DataLoader, DictDataLoader],
+) -> None:
+    # TODO: implement me
+    pass
+
+
 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
@@ -58,6 +70,32 @@ def plot_mlflow(
         writer.log_figure(fig, descr)
 
 
+def log_model_mlflow(
+    writer: Any,
+    model: Module,
+    dataloader: Union[None, DataLoader, DictDataLoader],
+) -> None:
+    if dataloader is not None:
+        xs = next(iter(dataloader))
+        if isinstance(xs, Sequence):
+            # ignore labels in supervised learning
+            xs = xs[0]
+        if isinstance(dataloader, DataLoader):
+            preds = model(xs)
+            xs = xs.numpy()
+            preds = preds.detach().numpy()
+        elif isinstance(dataloader, DictDataLoader):
+            preds = model(xs)
+            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=writer.get_artifact_uri(), signature=signature)
+
+
 TRACKER_MAPPING = {
     ExperimentTrackingTool.TENSORBOARD: write_tensorboard,
     ExperimentTrackingTool.MLFLOW: write_mlflow,
@@ -73,6 +111,11 @@ def plot_mlflow(
     ExperimentTrackingTool.MLFLOW: plot_mlflow,
 }
 
+MODEL_LOGGER_MAPPING = {
+    ExperimentTrackingTool.TENSORBOARD: log_model_tensorboard,
+    ExperimentTrackingTool.MLFLOW: log_model_mlflow,
+}
+
 
 def write_tracker(
     args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
@@ -90,3 +133,9 @@ def plot_tracker(
     args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
 ) -> None:
     return PLOTTER_MAPPING[tracking_tool](*args)  # type: ignore
+
+
+def log_model_tracker(
+    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return MODEL_LOGGER_MAPPING[tracking_tool](*args)  # type: ignore
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index 71963eab7..6e2352a18 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -23,6 +23,7 @@
 from qadence.ml_tools.data import DictDataLoader
 from qadence.ml_tools.optimize_step import optimize_step
 from qadence.ml_tools.printing import (
+    log_model_tracker,
     log_tracker,
     plot_tracker,
     print_metrics,
@@ -213,6 +214,9 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     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)
@@ -220,7 +224,6 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer.close()
     elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:
-        # track some final stuff e.g. images
         writer.end_run()
 
     return model, optimizer
diff --git a/qadence/ml_tools/train_no_grad.py b/qadence/ml_tools/train_no_grad.py
index 1b89dfbf6..02fea1119 100644
--- a/qadence/ml_tools/train_no_grad.py
+++ b/qadence/ml_tools/train_no_grad.py
@@ -21,7 +21,13 @@
 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 plot_tracker, print_metrics, write_tracker
+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
@@ -129,10 +135,19 @@ def _update_parameters(
             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_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+
+    # close tracker
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer.close()
     elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:

From 64fd8012000242e1a058608c2c0304c7d6c96451 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Mon, 15 Jul 2024 11:46:03 +0200
Subject: [PATCH 24/51] Add warning for tb model logger

---
 qadence/ml_tools/printing.py   | 6 ++++--
 qadence/ml_tools/train_grad.py | 2 ++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index be0d92a68..ce5c1ac55 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from logging import getLogger
 from typing import Any, Callable, Sequence, Union
 
 from mlflow.models import infer_signature
@@ -10,6 +11,8 @@
 from qadence.ml_tools.data import DictDataLoader
 from qadence.types import ExperimentTrackingTool
 
+logger = getLogger(__name__)
+
 
 def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:
     msg = " ".join(
@@ -46,8 +49,7 @@ def log_model_tensorboard(
     model: Module,
     dataloader: Union[None, DataLoader, DictDataLoader],
 ) -> None:
-    # TODO: implement me
-    pass
+    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:
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index 6e2352a18..f6e847a0a 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -221,6 +221,8 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     if config.folder:
         write_checkpoint(config.folder, model, optimizer, iteration)
     write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+
+    # close tracker
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
         writer.close()
     elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:

From 99f42be8c8b5bc8ecf5127038ff8720eff123d97 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Mon, 15 Jul 2024 13:24:21 +0200
Subject: [PATCH 25/51] Add plot_tensorboard

---
 qadence/ml_tools/printing.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index ce5c1ac55..59b1bd696 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -3,6 +3,7 @@
 from logging import getLogger
 from typing import Any, Callable, Sequence, Union
 
+from matplotlib.figure import Figure
 from mlflow.models import infer_signature
 from torch.nn import Module
 from torch.utils.data import DataLoader
@@ -13,6 +14,8 @@
 
 logger = getLogger(__name__)
 
+PlottingFunction = Callable[[Module, int], tuple[str, Figure]]
+
 
 def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:
     msg = " ".join(
@@ -38,10 +41,11 @@ def plot_tensorboard(
     writer: SummaryWriter,
     model: Module,
     iteration: int,
-    plotting_functions: tuple[Callable],
+    plotting_functions: tuple[PlottingFunction],
 ) -> None:
-    # TODO: implement me
-    pass
+    for pf in plotting_functions:
+        descr, fig = pf(model, iteration)
+        writer.add_figure(descr, fig, global_step=iteration)
 
 
 def log_model_tensorboard(
@@ -65,7 +69,7 @@ def plot_mlflow(
     writer: Any,
     model: Module,
     iteration: int,
-    plotting_functions: tuple[Callable],
+    plotting_functions: tuple[PlottingFunction],
 ) -> None:
     for pf in plotting_functions:
         descr, fig = pf(model, iteration)

From 8bbb293b557858c1765e21b3f176555283f7c6c9 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Mon, 15 Jul 2024 14:06:46 +0200
Subject: [PATCH 26/51] Correct log_model default and test_hp_log_mlflow

---
 qadence/ml_tools/config.py     |  2 +-
 tests/ml_tools/test_logging.py | 27 +++++++++++++++++++++++----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 13c48d02e..aebf99270 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -42,7 +42,7 @@ class TrainConfig:
 
     NOTE: currently only works with mlflow.
     """
-    log_model: bool = True
+    log_model: bool = False
     """Logs a serialised version of the model."""
     folder: Path | None = None
     """Checkpoint/tensorboard logs folder."""
diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
index f750b059a..74cf65d29 100644
--- a/tests/ml_tools/test_logging.py
+++ b/tests/ml_tools/test_logging.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import shutil
 from itertools import count
 from pathlib import Path
 
@@ -48,11 +49,29 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     experiment_id = mlflow_config.run.info.experiment_id
     run_id = mlflow_config.run.info.run_id
 
-    hyperparams_files = [
-        Path(f"mlruns/{experiment_id}/{run_id}/params/{key}") for key in hyperparams.keys()
-    ]
+    experiment_dir = Path(f"mlruns/{experiment_id}")
+    hyperparams_files = [experiment_dir / run_id / "params" / key for key in hyperparams.keys()]
+
     assert all([os.path.isfile(hf) for hf in hyperparams_files])
 
+    shutil.rmtree(experiment_dir)
+
+
+def test_hyperparams_logging_tensorboard() -> None:
+    pass
+
+
+def test_model_logging_mlflow() -> None:
+    pass
+
+
+def test_model_logging_tensorboard() -> None:
+    pass
+
+
+def test_plotting_mlflow() -> None:
+    pass
+
 
-def test_plotting() -> None:
+def test_plotting_tensorboard() -> None:
     pass

From 4dcfcdafae8a5a890b86e604ffe3a0bae056cd54 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Mon, 15 Jul 2024 17:54:29 +0200
Subject: [PATCH 27/51] Assign uuid to default mlflow run id

---
 qadence/ml_tools/config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index aebf99270..d4359f4b9 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -151,7 +151,7 @@ def __init__(self) -> None:
         If None or empty, a new experiment is created with a random UUID.
         """
 
-        self.run_name: str = os.getenv("MLFLOW_RUN_NAME", "test_0")
+        self.run_name: str = os.getenv("MLFLOW_RUN_NAME", str(uuid4()))
         """The name of the run."""
 
         if self.tracking_username != "":

From cabf464e833a8db21cdde5077a4fa68f8e179436 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Mon, 15 Jul 2024 17:58:00 +0200
Subject: [PATCH 28/51] Test model logging with tb/mlflow

---
 qadence/ml_tools/printing.py   |   5 +-
 tests/ml_tools/test_logging.py | 127 ++++++++++++++++++++++++++++++---
 2 files changed, 120 insertions(+), 12 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 59b1bd696..01319095f 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -86,12 +86,11 @@ def log_model_mlflow(
         if isinstance(xs, Sequence):
             # ignore labels in supervised learning
             xs = xs[0]
+        preds = model(xs)
         if isinstance(dataloader, DataLoader):
-            preds = model(xs)
             xs = xs.numpy()
             preds = preds.detach().numpy()
         elif isinstance(dataloader, DictDataLoader):
-            preds = model(xs)
             for key, val in xs.items():
                 xs[key] = val.numpy()
             for key, val in preds.items():
@@ -99,7 +98,7 @@ def log_model_mlflow(
         signature = infer_signature(xs, preds)
     else:
         signature = None
-    writer.pytorch.log_model(model, artifact_path=writer.get_artifact_uri(), signature=signature)
+    writer.pytorch.log_model(model, artifact_path="model", signature=signature)
 
 
 TRACKER_MAPPING = {
diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
index 74cf65d29..b4d9abe26 100644
--- a/tests/ml_tools/test_logging.py
+++ b/tests/ml_tools/test_logging.py
@@ -5,12 +5,15 @@
 from itertools import count
 from pathlib import Path
 
+import mlflow
+import pytest
 import torch
 from torch.utils.data import DataLoader
 
 from qadence.ml_tools import TrainConfig, train_with_grad
 from qadence.ml_tools.data import to_dataloader
-from qadence.models import QuantumModel
+from qadence.ml_tools.utils import rand_featureparameters
+from qadence.models import QNN, QuantumModel
 from qadence.types import ExperimentTrackingTool
 
 
@@ -21,7 +24,6 @@ def dataloader(batch_size: int = 25) -> DataLoader:
 
 
 def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
-    data = dataloader()
     model = BasicQuantumModel
     cnt = count()
     criterion = torch.nn.MSELoss()
@@ -43,7 +45,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
         tracking_tool=ExperimentTrackingTool.MLFLOW,
     )
 
-    train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+    train_with_grad(model, None, optimizer, config, loss_fn=loss_fn)
 
     mlflow_config = config.mlflow_config
     experiment_id = mlflow_config.run.info.experiment_id
@@ -57,16 +59,123 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     shutil.rmtree(experiment_dir)
 
 
-def test_hyperparams_logging_tensorboard() -> None:
-    pass
+def test_hyperparams_logging_tensorboard(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
+    model = BasicQuantumModel
+    cnt = count()
+    criterion = torch.nn.MSELoss()
+    hyperparams = {"max_iter": int(10), "lr": 0.1}
+    optimizer = torch.optim.Adam(model.parameters(), lr=hyperparams["lr"])
 
+    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+        next(cnt)
+        out = model.expectation({})
+        loss = criterion(out, torch.rand(1))
+        return loss, {}
 
-def test_model_logging_mlflow() -> None:
-    pass
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=hyperparams["max_iter"],  # type: ignore
+        checkpoint_every=1,
+        write_every=1,
+        hyperparams=hyperparams,
+        tracking_tool=ExperimentTrackingTool.TENSORBOARD,
+    )
 
+    train_with_grad(model, None, optimizer, config, loss_fn=loss_fn)
 
-def test_model_logging_tensorboard() -> None:
-    pass
+
+def test_model_logging_mlflow_basicQM(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
+    model = BasicQuantumModel
+    cnt = count()
+    criterion = torch.nn.MSELoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+
+    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+        next(cnt)
+        out = model.expectation({})
+        loss = criterion(out, torch.rand(1))
+        return loss, {}
+
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=10,  # type: ignore
+        checkpoint_every=1,
+        write_every=1,
+        log_model=True,
+        tracking_tool=ExperimentTrackingTool.MLFLOW,
+    )
+
+    train_with_grad(model, None, optimizer, config, loss_fn=loss_fn)
+
+    experiment_id = config.mlflow_config.run.info.experiment_id
+    run_id = config.mlflow_config.run.info.run_id
+
+    loaded_model = mlflow.pytorch.load_model(model_uri=f"runs:/{run_id}/model")
+
+    experiment_dir = Path(f"mlruns/{experiment_id}")
+    shutil.rmtree(experiment_dir)
+
+
+def test_model_logging_mlflow_basicQNN(BasicQNN: QNN, tmp_path: Path) -> None:
+    data = dataloader()
+    model = BasicQNN
+    cnt = count()
+    criterion = torch.nn.MSELoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+    inputs = rand_featureparameters(model, 1)
+
+    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+        next(cnt)
+        out = model.expectation(inputs)
+        loss = criterion(out, torch.rand(1))
+        return loss, {}
+
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=10,  # type: ignore
+        checkpoint_every=1,
+        write_every=1,
+        log_model=True,
+        tracking_tool=ExperimentTrackingTool.MLFLOW,
+    )
+
+    train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+    experiment_id = config.mlflow_config.run.info.experiment_id
+    run_id = config.mlflow_config.run.info.run_id
+
+    loaded_model = mlflow.pytorch.load_model(model_uri=f"runs:/{run_id}/model")
+
+    experiment_dir = Path(f"mlruns/{experiment_id}")
+    shutil.rmtree(experiment_dir)
+
+
+def test_model_logging_tensorboard(
+    BasicQuantumModel: QuantumModel, tmp_path: Path, caplog: pytest.LogCaptureFixture
+) -> None:
+    model = BasicQuantumModel
+    cnt = count()
+    criterion = torch.nn.MSELoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+
+    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+        next(cnt)
+        out = model.expectation({})
+        loss = criterion(out, torch.rand(1))
+        return loss, {}
+
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=10,  # type: ignore
+        checkpoint_every=1,
+        write_every=1,
+        log_model=True,
+        tracking_tool=ExperimentTrackingTool.TENSORBOARD,
+    )
+
+    train_with_grad(model, None, optimizer, config, loss_fn=loss_fn)
+
+    assert "Model logging is not supported by tensorboard. No model will be logged." in caplog.text
 
 
 def test_plotting_mlflow() -> None:

From 99e49ee50c733e9b164162f008da2e2f238bb7cd Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Mon, 15 Jul 2024 23:34:06 +0200
Subject: [PATCH 29/51] Refactor input data checking in log_model_mlflow

---
 qadence/ml_tools/printing.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 01319095f..9d14fa133 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 from logging import getLogger
-from typing import Any, Callable, Sequence, Union
+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
@@ -15,6 +16,7 @@
 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:
@@ -82,15 +84,13 @@ def log_model_mlflow(
     dataloader: Union[None, DataLoader, DictDataLoader],
 ) -> None:
     if dataloader is not None:
-        xs = next(iter(dataloader))
-        if isinstance(xs, Sequence):
-            # ignore labels in supervised learning
-            xs = xs[0]
+        xs: InputData
+        xs, *_ = next(iter(dataloader))
         preds = model(xs)
-        if isinstance(dataloader, DataLoader):
+        if isinstance(xs, Tensor):
             xs = xs.numpy()
             preds = preds.detach().numpy()
-        elif isinstance(dataloader, DictDataLoader):
+        elif isinstance(xs, dict):
             for key, val in xs.items():
                 xs[key] = val.numpy()
             for key, val in preds.items():

From 4a12155b5caa16eff38b97ba38bf087a38fc3cb4 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Tue, 16 Jul 2024 08:28:11 +0200
Subject: [PATCH 30/51] Refactor mlflow tests for clarity

---
 tests/ml_tools/test_logging.py | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
index b4d9abe26..d5d498bb3 100644
--- a/tests/ml_tools/test_logging.py
+++ b/tests/ml_tools/test_logging.py
@@ -23,6 +23,18 @@ def dataloader(batch_size: int = 25) -> DataLoader:
     return to_dataloader(x, y, batch_size=batch_size, infinite=True)
 
 
+def load_mlflow_model(train_config: TrainConfig) -> None:
+    run_id = train_config.mlflow_config.run.info.run_id
+
+    mlflow.pytorch.load_model(model_uri=f"runs:/{run_id}/model")
+
+
+def clean_mlflow_experiment_dir(train_config: TrainConfig) -> None:
+    experiment_id = train_config.mlflow_config.run.info.experiment_id
+    experiment_dir = Path(f"mlruns/{experiment_id}")
+    shutil.rmtree(experiment_dir)
+
+
 def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
     model = BasicQuantumModel
     cnt = count()
@@ -56,7 +68,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
 
     assert all([os.path.isfile(hf) for hf in hyperparams_files])
 
-    shutil.rmtree(experiment_dir)
+    clean_mlflow_experiment_dir(config)
 
 
 def test_hyperparams_logging_tensorboard(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
@@ -107,13 +119,9 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
 
     train_with_grad(model, None, optimizer, config, loss_fn=loss_fn)
 
-    experiment_id = config.mlflow_config.run.info.experiment_id
-    run_id = config.mlflow_config.run.info.run_id
-
-    loaded_model = mlflow.pytorch.load_model(model_uri=f"runs:/{run_id}/model")
+    load_mlflow_model(config)
 
-    experiment_dir = Path(f"mlruns/{experiment_id}")
-    shutil.rmtree(experiment_dir)
+    clean_mlflow_experiment_dir(config)
 
 
 def test_model_logging_mlflow_basicQNN(BasicQNN: QNN, tmp_path: Path) -> None:
@@ -141,13 +149,9 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
 
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
 
-    experiment_id = config.mlflow_config.run.info.experiment_id
-    run_id = config.mlflow_config.run.info.run_id
+    load_mlflow_model(config)
 
-    loaded_model = mlflow.pytorch.load_model(model_uri=f"runs:/{run_id}/model")
-
-    experiment_dir = Path(f"mlruns/{experiment_id}")
-    shutil.rmtree(experiment_dir)
+    clean_mlflow_experiment_dir(config)
 
 
 def test_model_logging_tensorboard(

From 8691327a05134d789dee9a4b87909eee95ca60aa Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Tue, 16 Jul 2024 09:46:13 +0200
Subject: [PATCH 31/51] Refactor tests and test adjoint model log

---
 tests/ml_tools/test_logging.py | 94 ++++++++++++++++++----------------
 1 file changed, 50 insertions(+), 44 deletions(-)

diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
index d5d498bb3..2409ea0a3 100644
--- a/tests/ml_tools/test_logging.py
+++ b/tests/ml_tools/test_logging.py
@@ -4,10 +4,13 @@
 import shutil
 from itertools import count
 from pathlib import Path
+from typing import Callable
 
 import mlflow
 import pytest
 import torch
+from torch.nn import Module
+from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 
 from qadence.ml_tools import TrainConfig, train_with_grad
@@ -23,6 +26,21 @@ def dataloader(batch_size: int = 25) -> DataLoader:
     return to_dataloader(x, y, batch_size=batch_size, infinite=True)
 
 
+def setup(model: Module) -> tuple[Callable, Optimizer]:
+    cnt = count()
+    criterion = torch.nn.MSELoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+    inputs = rand_featureparameters(model, 1)
+
+    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+        next(cnt)
+        out = model.expectation(inputs)
+        loss = criterion(out, torch.rand(1))
+        return loss, {}
+
+    return loss_fn, optimizer
+
+
 def load_mlflow_model(train_config: TrainConfig) -> None:
     run_id = train_config.mlflow_config.run.info.run_id
 
@@ -37,16 +55,10 @@ def clean_mlflow_experiment_dir(train_config: TrainConfig) -> None:
 
 def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
     model = BasicQuantumModel
-    cnt = count()
-    criterion = torch.nn.MSELoss()
-    hyperparams = {"max_iter": int(10), "lr": 0.1}
-    optimizer = torch.optim.Adam(model.parameters(), lr=hyperparams["lr"])
 
-    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
-        next(cnt)
-        out = model.expectation({})
-        loss = criterion(out, torch.rand(1))
-        return loss, {}
+    loss_fn, optimizer = setup(model)
+
+    hyperparams = {"max_iter": int(10), "lr": 0.1}
 
     config = TrainConfig(
         folder=tmp_path,
@@ -73,16 +85,10 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
 
 def test_hyperparams_logging_tensorboard(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
     model = BasicQuantumModel
-    cnt = count()
-    criterion = torch.nn.MSELoss()
-    hyperparams = {"max_iter": int(10), "lr": 0.1}
-    optimizer = torch.optim.Adam(model.parameters(), lr=hyperparams["lr"])
 
-    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
-        next(cnt)
-        out = model.expectation({})
-        loss = criterion(out, torch.rand(1))
-        return loss, {}
+    loss_fn, optimizer = setup(model)
+
+    hyperparams = {"max_iter": int(10), "lr": 0.1}
 
     config = TrainConfig(
         folder=tmp_path,
@@ -98,15 +104,8 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
 
 def test_model_logging_mlflow_basicQM(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
     model = BasicQuantumModel
-    cnt = count()
-    criterion = torch.nn.MSELoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
 
-    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
-        next(cnt)
-        out = model.expectation({})
-        loss = criterion(out, torch.rand(1))
-        return loss, {}
+    loss_fn, optimizer = setup(model)
 
     config = TrainConfig(
         folder=tmp_path,
@@ -127,16 +126,30 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
 def test_model_logging_mlflow_basicQNN(BasicQNN: QNN, tmp_path: Path) -> None:
     data = dataloader()
     model = BasicQNN
-    cnt = count()
-    criterion = torch.nn.MSELoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
-    inputs = rand_featureparameters(model, 1)
 
-    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
-        next(cnt)
-        out = model.expectation(inputs)
-        loss = criterion(out, torch.rand(1))
-        return loss, {}
+    loss_fn, optimizer = setup(model)
+
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=10,  # type: ignore
+        checkpoint_every=1,
+        write_every=1,
+        log_model=True,
+        tracking_tool=ExperimentTrackingTool.MLFLOW,
+    )
+
+    train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+    load_mlflow_model(config)
+
+    clean_mlflow_experiment_dir(config)
+
+
+def test_model_logging_mlflow_basicAdjQNN(BasicAdjointQNN: QNN, tmp_path: Path) -> None:
+    data = dataloader()
+    model = BasicAdjointQNN
+
+    loss_fn, optimizer = setup(model)
 
     config = TrainConfig(
         folder=tmp_path,
@@ -158,15 +171,8 @@ def test_model_logging_tensorboard(
     BasicQuantumModel: QuantumModel, tmp_path: Path, caplog: pytest.LogCaptureFixture
 ) -> None:
     model = BasicQuantumModel
-    cnt = count()
-    criterion = torch.nn.MSELoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
 
-    def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
-        next(cnt)
-        out = model.expectation({})
-        loss = criterion(out, torch.rand(1))
-        return loss, {}
+    loss_fn, optimizer = setup(model)
 
     config = TrainConfig(
         folder=tmp_path,
@@ -182,7 +188,7 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     assert "Model logging is not supported by tensorboard. No model will be logged." in caplog.text
 
 
-def test_plotting_mlflow() -> None:
+def test_plotting_mlflow(BasicQNN: QNN, tmp_path: Path) -> None:
     pass
 
 

From 33f857cb98475e1d6df3008a1132f795080472db Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Tue, 16 Jul 2024 14:20:49 +0200
Subject: [PATCH 32/51] Add plotting test and refactor typing

Add LoggablePlotFunction type and correct plotting_functions type in
TrainConfig
---
 qadence/ml_tools/config.py     |   4 +-
 qadence/types.py               |   5 ++
 tests/ml_tools/test_logging.py | 108 ++++++++++++++++++++++++++++++---
 3 files changed, 105 insertions(+), 12 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index d4359f4b9..ca5854095 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -10,7 +10,7 @@
 
 from torch import Tensor
 
-from qadence.types import ExperimentTrackingTool
+from qadence.types import ExperimentTrackingTool, LoggablePlotFunction
 
 logger = getLogger(__name__)
 
@@ -65,7 +65,7 @@ class TrainConfig:
     """The tracking tool of choice."""
     hyperparams: dict = field(default_factory=dict)
     """Hyperparameters to track."""
-    plotting_functions: tuple[Callable] = field(default_factory=tuple)  # type: ignore
+    plotting_functions: tuple[LoggablePlotFunction, ...] = field(default_factory=tuple)  # type: ignore
     """Functions for in-train plotting."""
 
     # tensorboard only allows for certain types as hyperparameters
diff --git a/qadence/types.py b/qadence/types.py
index 64f22352c..7a5a494fb 100644
--- a/qadence/types.py
+++ b/qadence/types.py
@@ -6,8 +6,10 @@
 
 import numpy as np
 import sympy
+from matplotlib.figure import Figure
 from numpy.typing import ArrayLike
 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."""
@@ -408,3 +410,6 @@ class ExperimentTrackingTool(StrEnum):
     """Use the tensorboard experiment tracker."""
     MLFLOW = "mlflow"
     """Use the ml-flow experiment tracker."""
+
+
+LoggablePlotFunction = Callable[[Module, int], tuple[str, Figure]]
diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
index 2409ea0a3..afaa47e60 100644
--- a/tests/ml_tools/test_logging.py
+++ b/tests/ml_tools/test_logging.py
@@ -5,10 +5,15 @@
 from itertools import count
 from pathlib import Path
 from typing import Callable
+from urllib.parse import urlparse
 
 import mlflow
 import pytest
 import torch
+from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
+from mlflow import MlflowClient
+from mlflow.entities import Run
 from torch.nn import Module
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
@@ -47,10 +52,27 @@ def load_mlflow_model(train_config: TrainConfig) -> None:
     mlflow.pytorch.load_model(model_uri=f"runs:/{run_id}/model")
 
 
-def clean_mlflow_experiment_dir(train_config: TrainConfig) -> None:
+def clean_mlflow_experiment(train_config: TrainConfig) -> None:
     experiment_id = train_config.mlflow_config.run.info.experiment_id
-    experiment_dir = Path(f"mlruns/{experiment_id}")
-    shutil.rmtree(experiment_dir)
+    client = MlflowClient()
+
+    runs = client.search_runs(experiment_id)
+
+    def clean_artifacts(run: Run) -> None:
+        artifact_uri = run.info.artifact_uri
+        parsed_uri = urlparse(artifact_uri)
+        local_path = os.path.abspath(os.path.join(parsed_uri.netloc, parsed_uri.path))
+        shutil.rmtree(local_path)
+
+    for run in runs:
+        clean_artifacts(run)
+
+        run_id = run.info.run_id
+        client.delete_run(run_id)
+
+        mlruns_base_dir = "./mlruns"
+        if os.path.isdir(mlruns_base_dir):
+            shutil.rmtree(os.path.join(mlruns_base_dir, experiment_id))
 
 
 def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
@@ -80,7 +102,7 @@ def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: P
 
     assert all([os.path.isfile(hf) for hf in hyperparams_files])
 
-    clean_mlflow_experiment_dir(config)
+    clean_mlflow_experiment(config)
 
 
 def test_hyperparams_logging_tensorboard(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None:
@@ -120,7 +142,7 @@ def test_model_logging_mlflow_basicQM(BasicQuantumModel: QuantumModel, tmp_path:
 
     load_mlflow_model(config)
 
-    clean_mlflow_experiment_dir(config)
+    clean_mlflow_experiment(config)
 
 
 def test_model_logging_mlflow_basicQNN(BasicQNN: QNN, tmp_path: Path) -> None:
@@ -142,7 +164,7 @@ def test_model_logging_mlflow_basicQNN(BasicQNN: QNN, tmp_path: Path) -> None:
 
     load_mlflow_model(config)
 
-    clean_mlflow_experiment_dir(config)
+    clean_mlflow_experiment(config)
 
 
 def test_model_logging_mlflow_basicAdjQNN(BasicAdjointQNN: QNN, tmp_path: Path) -> None:
@@ -164,7 +186,7 @@ def test_model_logging_mlflow_basicAdjQNN(BasicAdjointQNN: QNN, tmp_path: Path)
 
     load_mlflow_model(config)
 
-    clean_mlflow_experiment_dir(config)
+    clean_mlflow_experiment(config)
 
 
 def test_model_logging_tensorboard(
@@ -189,8 +211,74 @@ def test_model_logging_tensorboard(
 
 
 def test_plotting_mlflow(BasicQNN: QNN, tmp_path: Path) -> None:
-    pass
+    data = dataloader()
+    model = BasicQNN
+
+    loss_fn, optimizer = setup(model)
 
+    def plot_model(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
+        descr = f"model_prediction_epoch_{iteration}.png"
+        fig, ax = plt.subplots()
+        x = torch.linspace(0, 1, 100).reshape(-1, 1)
+        out = model.expectation(x)
+        ax.plot(x.detach().numpy(), out.detach().numpy())
+        return descr, fig
+
+    def plot_error(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
+        descr = f"error_epoch_{iteration}.png"
+        fig, ax = plt.subplots()
+        x = torch.linspace(0, 1, 100).reshape(-1, 1)
+        out = model.expectation(x)
+        ground_truth = torch.rand_like(out)
+        error = ground_truth - out
+        ax.plot(x.detach().numpy(), error.detach().numpy())
+        return descr, fig
 
-def test_plotting_tensorboard() -> None:
-    pass
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=10,
+        checkpoint_every=1,
+        write_every=1,
+        tracking_tool=ExperimentTrackingTool.MLFLOW,
+        plotting_functions=(plot_model, plot_error),
+    )
+
+    train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+    clean_mlflow_experiment(config)
+
+
+def test_plotting_tensorboard(BasicQNN: QNN, tmp_path: Path) -> None:
+    data = dataloader()
+    model = BasicQNN
+
+    loss_fn, optimizer = setup(model)
+
+    def plot_model(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
+        descr = f"model_prediction_epoch_{iteration}.png"
+        fig, ax = plt.subplots()
+        x = torch.linspace(0, 1, 100).reshape(-1, 1)
+        out = model.expectation(x)
+        ax.plot(x.detach().numpy(), out.detach().numpy())
+        return descr, fig
+
+    def plot_error(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
+        descr = f"error_epoch_{iteration}.png"
+        fig, ax = plt.subplots()
+        x = torch.linspace(0, 1, 100).reshape(-1, 1)
+        out = model.expectation(x)
+        ground_truth = torch.rand_like(out)
+        error = ground_truth - out
+        ax.plot(x.detach().numpy(), error.detach().numpy())
+        return descr, fig
+
+    config = TrainConfig(
+        folder=tmp_path,
+        max_iter=10,
+        checkpoint_every=1,
+        write_every=1,
+        tracking_tool=ExperimentTrackingTool.TENSORBOARD,
+        plotting_functions=(plot_model, plot_error),
+    )
+
+    train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)

From 13a5ea191d9de26efd2dcf3d9b428aac0010f864 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Tue, 16 Jul 2024 16:30:14 +0200
Subject: [PATCH 33/51] Test to see if img files are produced

---
 tests/ml_tools/test_logging.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
index afaa47e60..8b779e08c 100644
--- a/tests/ml_tools/test_logging.py
+++ b/tests/ml_tools/test_logging.py
@@ -52,6 +52,12 @@ def load_mlflow_model(train_config: TrainConfig) -> None:
     mlflow.pytorch.load_model(model_uri=f"runs:/{run_id}/model")
 
 
+def find_mlflow_artifacts_path(run: Run) -> Path:
+    artifact_uri = run.info.artifact_uri
+    parsed_uri = urlparse(artifact_uri)
+    return Path(os.path.abspath(os.path.join(parsed_uri.netloc, parsed_uri.path)))
+
+
 def clean_mlflow_experiment(train_config: TrainConfig) -> None:
     experiment_id = train_config.mlflow_config.run.info.experiment_id
     client = MlflowClient()
@@ -59,9 +65,7 @@ def clean_mlflow_experiment(train_config: TrainConfig) -> None:
     runs = client.search_runs(experiment_id)
 
     def clean_artifacts(run: Run) -> None:
-        artifact_uri = run.info.artifact_uri
-        parsed_uri = urlparse(artifact_uri)
-        local_path = os.path.abspath(os.path.join(parsed_uri.netloc, parsed_uri.path))
+        local_path = find_mlflow_artifacts_path(run)
         shutil.rmtree(local_path)
 
     for run in runs:
@@ -234,17 +238,27 @@ def plot_error(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
         ax.plot(x.detach().numpy(), error.detach().numpy())
         return descr, fig
 
+    max_iter = 10
+    plot_every = 2
     config = TrainConfig(
         folder=tmp_path,
-        max_iter=10,
+        max_iter=max_iter,
         checkpoint_every=1,
         write_every=1,
+        plot_every=plot_every,
         tracking_tool=ExperimentTrackingTool.MLFLOW,
         plotting_functions=(plot_model, plot_error),
     )
 
     train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
 
+    all_plot_names = [f"model_prediction_epoch_{i}.png" for i in range(0, max_iter, plot_every)]
+    all_plot_names.extend([f"error_epoch_{i}.png" for i in range(0, max_iter, plot_every)])
+
+    artifact_path = find_mlflow_artifacts_path(config.mlflow_config.run)
+
+    assert all([os.path.isfile(artifact_path / pn) for pn in all_plot_names])
+
     clean_mlflow_experiment(config)
 
 

From 3d26ab70156977f04b8dd44aebf199ed23b98fd6 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Tue, 16 Jul 2024 21:07:31 +0200
Subject: [PATCH 34/51] Remove user/pwd from MLConfig

Handled at tracking server setup level
---
 qadence/ml_tools/config.py | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index ca5854095..d2960ce83 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -125,8 +125,8 @@ class MLFlowConfig:
     Example:
 
         export MLFLOW_TRACKING_URI=tracking_uri
-        export MLFLOW_TRACKING_USERNAME=username
-        export MLFLOW_TRACKING_PASSWORD=password
+        export MLFLOW_EXPERIMENT=experiment_name
+        export MLFLOW_RUN_NAME=run_name
     """
 
     def __init__(self) -> None:
@@ -139,12 +139,6 @@ def __init__(self) -> None:
         Data is stored locally at the provided file (or ./mlruns if empty).
         """
 
-        self.tracking_username: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
-        """The username for the mlflow tracking server."""
-
-        self.tracking_pwd: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
-        """The password for the mlflow tracking server."""
-
         self.experiment_name: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
         """The name of the experiment.
 
@@ -154,10 +148,8 @@ def __init__(self) -> None:
         self.run_name: str = os.getenv("MLFLOW_RUN_NAME", str(uuid4()))
         """The name of the run."""
 
-        if self.tracking_username != "":
-            logger.info(f"Intialized mlflow remote logging for user {self.tracking_username}.")
-
         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):

From 21245b6598283e13320a18adbf5d62ef2341bc7a Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 11:47:56 +0200
Subject: [PATCH 35/51] Add mlflow tutorial in docs

---
 docs/tutorials/qml/ml_tools.md             | 113 +++++++++++++++++----
 docs/tutorials/qml/mlflow_demonstration.py |  91 -----------------
 2 files changed, 91 insertions(+), 113 deletions(-)
 delete mode 100644 docs/tutorials/qml/mlflow_demonstration.py

diff --git a/docs/tutorials/qml/ml_tools.md b/docs/tutorials/qml/ml_tools.md
index c5f40fdaf..00154c4c2 100644
--- a/docs/tutorials/qml/ml_tools.md
+++ b/docs/tutorials/qml/ml_tools.md
@@ -285,46 +285,85 @@ def train(
     return model, optimizer
 ```
 
-## MLTools with mlflow
+## Experiment tracking with mlflow
 
-MLTools now offers MLflow support.
+Qadence allows to track runs and log hyperparameters, models and plots with [tensorboard](https://pytorch.org/tutorials/recipes/recipes/tensorboard_with_pytorch.html) and [mlflow](https://mlflow.org/). In the following, we demonstrate the integration with mlflow.
 
+### mlflow configuration
+We have control over our tracking configuration by setting environment variables. First, let's look at the trakcing URI. For the purpose of this demo we will be working with a local database, in a similar fashion as described [here](https://mlflow.org/docs/latest/tracking/tutorials/local-database.html),
+```bash
+export MLFLOW_TRACKING_URI=sqlite:///mlruns.db
+```
 
-```python
-from __future__ import annotations
+Qadence can also read the following two environment variables to define the mlflow experiment name and run name
+```bash
+export MLFLOW_EXPERIMENT=test_experiment
+export MLFLOW_RUN_NAME=run_0
+```
+
+If no tracking URI is provided, mlflow stores runs and artifacts in the local `./mlflow` directory and if no names are defined, the experiment and run will be named with random UUIDs.
 
-import os
+### Setup
+Let's do the necessary imports and declare a `DataLoader`. We can already define some hyperparameters here, including the seed for random number generators. mlflow can log hyperparameters with arbitrary types, for example the observable that we want to monitor (`Z` in this case, which has a `qadence.Operation` type).
+
+```python
+import random
 from itertools import count
+
+import numpy as np
 import torch
+from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
+from torch.nn import Module
 from torch.utils.data import DataLoader
 
-from qadence.ml_tools import (
-    TrainConfig,
-    train_with_grad,
-
-)
+from qadence import hea, QuantumCircuit, Z
+from qadence.constructors import feature_map, hamiltonian_factory
+from qadence.ml_tools import train_with_grad, TrainConfig
 from qadence.ml_tools.data import to_dataloader
 from qadence.ml_tools.utils import rand_featureparameters
 from qadence.models import QNN, QuantumModel
 from qadence.types import ExperimentTrackingTool
-from qadence import QuantumCircuit, hea, Z
 
-os.environ['MLFLOW_TRACKING_URI'] = 'sqlite:///mlflow.db'
-os.environ['MLFLOW_EXPERIMENT'] = 'mlflow_demonstration'
-os.environ['MLFLOW_RUN_NAME'] = 'test_0'
+hyperparams = {
+    "seed": 42,
+    "batch_size": 10,
+    "n_qubits": 2,
+    "ansatz_depth": 1,
+    "observable": Z,
+}
+
+np.random.seed(hyperparams["seed"])
+torch.manual_seed(hyperparams["seed"])
+random.seed(hyperparams["seed"])
+
 
-# in case you want to track remotely
-#os.environ['MLFLOW_TRACKING_USERNAME'] =
-#s.environ['MLFLOW_TRACKING_PASSWORD'] =
 def dataloader(batch_size: int = 25) -> DataLoader:
     x = torch.linspace(0, 1, batch_size).reshape(-1, 1)
     y = torch.cos(x)
     return to_dataloader(x, y, batch_size=batch_size, infinite=True)
-data = dataloader()
-model = QNN(QuantumCircuit(2, hea(2,1)), observable=Z(0))
+```
+
+We continue with the usual QNN definition, as well as loss function and optimizer.
+
+```python
+obs = hamiltonian_factory(register=hyperparams["n_qubits"], detuning=hyperparams["observable"])
+
+data = dataloader(hyperparams["batch_size"])
+fm = feature_map(hyperparams["n_qubits"], param="x")
+
+model = QNN(
+    QuantumCircuit(
+        hyperparams["n_qubits"], fm, hea(hyperparams["n_qubits"], hyperparams["ansatz_depth"])
+    ),
+    observable=obs,
+    inputs=["x"],
+)
+
 cnt = count()
 criterion = torch.nn.MSELoss()
 optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+
 inputs = rand_featureparameters(model, 1)
 
 def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
@@ -332,12 +371,42 @@ def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict
     out = model.expectation(inputs)
     loss = criterion(out, torch.rand(1))
     return loss, {}
+```
+
+### `TrainConfig` specifications
+Qadence offers different tracking options via `TrainConfig`. Here we use the `ExperimentTrackingTool` type to specify that we want to track the experiment with mlflow. Tracking with tensorboard is also possible. We can then indicate *what* and *how often* we want to track or log. `write_every` controls the number of epochs after which the loss values is logged. Thanks to the `plotting_functions` and `plot_every`arguments, we are also able to plot model-related quantities throughout training. Notice that arbitrary plotting functions can be passed, as long as the signature is the same as `plot_fn` below. Finally, the trained model can be logged by setting `log_model=True`. Here is an example of plotting function and training configuration
+
+```python
+def plot_fn(model: Module, iteration: int) -> tuple[str, Figure]:
+    descr = f"ufa_prediction_epoch_{iteration}.png"
+    fig, ax = plt.subplots()
+    x = torch.linspace(0, 1, 100).reshape(-1, 1)
+    out = model.expectation(x)
+    ax.plot(x.detach().numpy(), out.detach().numpy())
+    return descr, fig
+
 
 config = TrainConfig(
-    folder='mlflow_demonstration', max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=ExperimentTrackingTool.MLFLOW
+    folder="mlflow_demonstration",
+    max_iter=10,
+    checkpoint_every=1,
+    plot_every=2,
+    write_every=1,
+    log_model=True,
+    tracking_tool=ExperimentTrackingTool.MLFLOW,
+    hyperparams=hyperparams,
+    plotting_functions=(plot_fn,),
 )
+```
+
+### Training and inspecting
+Model training happens as usual
+```python
 train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+```
 
-os.system('mlflow ui --port 5000')
-os.system('mlflow ui --backend-store-uri sqlite:///mlflow.db')
+After training , we can inspect our experiment via the mlflow UI
+```bash
+mlflow ui --port 8080 --backend-store-uri sqlite:///mlruns.db
 ```
+In this case, since we're running on a local server, we can access the mlflow UI by navigating to http://localhost:8080/.
diff --git a/docs/tutorials/qml/mlflow_demonstration.py b/docs/tutorials/qml/mlflow_demonstration.py
deleted file mode 100644
index f34a94d72..000000000
--- a/docs/tutorials/qml/mlflow_demonstration.py
+++ /dev/null
@@ -1,91 +0,0 @@
-from __future__ import annotations
-
-import random
-from itertools import count
-
-import numpy as np
-import torch
-
-# from dotenv import load_dotenv
-from matplotlib import pyplot as plt
-from matplotlib.figure import Figure
-from torch.utils.data import DataLoader
-
-from qadence import QuantumCircuit, Z, hea
-from qadence.constructors import feature_map, hamiltonian_factory
-from qadence.ml_tools import TrainConfig, train_with_grad
-from qadence.ml_tools.data import to_dataloader
-from qadence.ml_tools.utils import rand_featureparameters
-from qadence.models import QNN, QuantumModel
-from qadence.types import ExperimentTrackingTool
-
-hyperparams = {
-    "seed": 42,
-    "batch_size": 10,
-    "n_qubits": 2,
-    "ansatz_depth": 1,
-    "observable": Z,
-}
-
-np.random.seed(hyperparams["seed"])
-torch.manual_seed(hyperparams["seed"])
-random.seed(hyperparams["seed"])
-
-# load_dotenv()
-
-
-# in case you want to track remotely
-# os.environ['MLFLOW_TRACKING_USERNAME'] =
-# s.environ['MLFLOW_TRACKING_PASSWORD'] =
-def dataloader(batch_size: int = 25) -> DataLoader:
-    x = torch.linspace(0, 1, batch_size).reshape(-1, 1)
-    y = torch.cos(x)
-    return to_dataloader(x, y, batch_size=batch_size, infinite=True)
-
-
-obs = hamiltonian_factory(register=hyperparams["n_qubits"], detuning=hyperparams["observable"])
-
-data = dataloader(hyperparams["batch_size"])
-fm = feature_map(hyperparams["n_qubits"], param="x")
-model = QNN(
-    QuantumCircuit(
-        hyperparams["n_qubits"], fm, hea(hyperparams["n_qubits"], hyperparams["ansatz_depth"])
-    ),
-    observable=obs,
-    inputs=["x"],
-)
-cnt = count()
-criterion = torch.nn.MSELoss()
-optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
-inputs = rand_featureparameters(model, 1)
-
-
-def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
-    next(cnt)
-    out = model.expectation(inputs)
-    loss = criterion(out, torch.rand(1))
-    return loss, {}
-
-
-def plot_fn(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
-    descr = f"ufa_prediction_epoch_{iteration}.png"
-    fig, ax = plt.subplots()
-    x = torch.linspace(0, 1, 100).reshape(-1, 1)
-    out = model.expectation(x)
-    ax.plot(x.detach().numpy(), out.detach().numpy())
-    return descr, fig
-
-
-config = TrainConfig(
-    folder="mlflow_demonstration",
-    max_iter=10,
-    checkpoint_every=1,
-    plot_every=2,
-    write_every=1,
-    log_model=True,
-    tracking_tool=ExperimentTrackingTool.MLFLOW,
-    hyperparams=hyperparams,
-    plotting_functions=(plot_fn,),
-)
-
-train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)

From 8e7fbae71bbc0ad48947005afac5079288de47bf Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 16:04:47 +0200
Subject: [PATCH 36/51] Fix mlflow dependency in pyproject

---
 pyproject.toml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 12d6f6aa2..89a69af8e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -84,8 +84,7 @@ protocols = ["qadence-protocols"]
 libs = ["qadence-libs"]
 dlprof = ["nvidia-pyindex", "nvidia-dlprof[pytorch]"]
 mlflow = ["mlflow"]
-all = ["pulser", "braket", "visualization", "protocols", "libs"]
-
+all = ["pulser", "braket", "visualization", "protocols", "libs", "mlflow"]
 
 [tool.hatch.envs.default]
 dependencies = [
@@ -103,7 +102,7 @@ dependencies = [
     "ruff",
     "pydocstringformatter",
 ]
-features = ["pulser", "braket", "visualization", "horqrux"]
+features = ["pulser", "braket", "visualization", "horqrux", "mlflow"]
 
 [tool.hatch.envs.default.scripts]
 test = "pytest -n auto --cov-report lcov --cov-config=pyproject.toml --cov=qadence --cov=tests --ignore=./tests/test_examples.py {args}"

From 5b7ddf002f814976ed742df357826a9cf9e43aea Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 16:18:14 +0200
Subject: [PATCH 37/51] Add author

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 89a69af8e..fcead79d7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,6 +21,7 @@ authors = [
     { name = "Smit Chaudhary", email = "smit.chaudhary@pasqal.com" },
     { name = "Ignacio Fernández Graña", email = "ignacio.fernandez-grana@pasqal.com" },
     { name = "Charles Moussa", email = "charles.moussa@pasqal.com" },
+    { name = "Giorgio Tosti Balducci", email = "giorgio.tosti-balducci@pasqal.com" },
 ]
 requires-python = ">=3.9"
 license = { text = "Apache 2.0" }

From 89ea8a135a2d82fc7a4ac2ead1c54b907ceef75f Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 16:24:30 +0200
Subject: [PATCH 38/51] Fix tests

---
 qadence/ml_tools/train_grad.py | 4 ++++
 tests/ml_tools/test_logging.py | 3 ++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index 2fb080a4d..d04265623 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -196,6 +196,10 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                 else:
                     write_checkpoint(config.folder, model, optimizer, init_iter)
 
+            plot_tracker(
+                (writer, model, init_iter, config.plotting_functions), config.tracking_tool
+            )
+
         except KeyboardInterrupt:
             logger.info("Terminating training gracefully after the current iteration.")
 
diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py
index 8b779e08c..d8669341e 100644
--- a/tests/ml_tools/test_logging.py
+++ b/tests/ml_tools/test_logging.py
@@ -20,8 +20,9 @@
 
 from qadence.ml_tools import TrainConfig, train_with_grad
 from qadence.ml_tools.data import to_dataloader
+from qadence.ml_tools.models import QNN
 from qadence.ml_tools.utils import rand_featureparameters
-from qadence.models import QNN, QuantumModel
+from qadence.model import QuantumModel
 from qadence.types import ExperimentTrackingTool
 
 

From a0ace68719e974af852ee84d1cd9960a927126dd Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 16:29:54 +0200
Subject: [PATCH 39/51] Bump version

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index fcead79d7..fbbf1207c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,7 @@ authors = [
 ]
 requires-python = ">=3.9"
 license = { text = "Apache 2.0" }
-version = "1.7.2"
+version = "1.7.3"
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python",

From a13cbd451a20fc0d1a899e4cabc7a7cc6f94ac8f Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 16:41:06 +0200
Subject: [PATCH 40/51] Fix docs

---
 pyproject.toml                 | 2 +-
 qadence/ml_tools/train_grad.py | 4 ----
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index fbbf1207c..a1368d006 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -140,7 +140,7 @@ dependencies = [
     "markdown-exec",
     "mike",
 ]
-features = ["pulser", "braket", "horqrux", "visualization"]
+features = ["pulser", "braket", "horqrux", "visualization", "mlflow"]
 
 [tool.hatch.envs.docs.scripts]
 build = "mkdocs build --clean --strict"
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index d04265623..a586baeaa 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -66,10 +66,6 @@ def train(
         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:

From f95bf198f71570ac9f3f73f294803a26b944b9d7 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 16:51:05 +0200
Subject: [PATCH 41/51] Fix TrainConfig docstring

---
 qadence/ml_tools/config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 286e48481..10ecf9b26 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -47,7 +47,7 @@ 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."""
     plot_every: int = 5000

From 14373c7c82ec5404e7e05c1da8cafc7f95dac58f Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 22:12:20 +0200
Subject: [PATCH 42/51] Correct typing in TrainConfig

---
 qadence/ml_tools/config.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py
index 10ecf9b26..d15ec78b4 100644
--- a/qadence/ml_tools/config.py
+++ b/qadence/ml_tools/config.py
@@ -5,7 +5,7 @@
 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
@@ -73,9 +73,9 @@ 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."""

From b018e0ed7c840491a59f5d943f33dced1f531f61 Mon Sep 17 00:00:00 2001
From: Giorgio Tosti Balducci
 <64412197+debrevitatevitae@users.noreply.github.com>
Date: Wed, 17 Jul 2024 22:17:01 +0200
Subject: [PATCH 43/51] Correct typo in docs/tutorials/qml/ml_tools.md

Co-authored-by: Roland-djee <9250798+Roland-djee@users.noreply.github.com>
---
 docs/tutorials/qml/ml_tools.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/tutorials/qml/ml_tools.md b/docs/tutorials/qml/ml_tools.md
index 769fcc04c..c53f09e5c 100644
--- a/docs/tutorials/qml/ml_tools.md
+++ b/docs/tutorials/qml/ml_tools.md
@@ -312,7 +312,7 @@ def train(
 Qadence allows to track runs and log hyperparameters, models and plots with [tensorboard](https://pytorch.org/tutorials/recipes/recipes/tensorboard_with_pytorch.html) and [mlflow](https://mlflow.org/). In the following, we demonstrate the integration with mlflow.
 
 ### mlflow configuration
-We have control over our tracking configuration by setting environment variables. First, let's look at the trakcing URI. For the purpose of this demo we will be working with a local database, in a similar fashion as described [here](https://mlflow.org/docs/latest/tracking/tutorials/local-database.html),
+We have control over our tracking configuration by setting environment variables. First, let's look at the tracking URI. For the purpose of this demo we will be working with a local database, in a similar fashion as described [here](https://mlflow.org/docs/latest/tracking/tutorials/local-database.html),
 ```bash
 export MLFLOW_TRACKING_URI=sqlite:///mlruns.db
 ```

From 243366ce62562709f4029293eb3af4e09fa2358e Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 22:19:22 +0200
Subject: [PATCH 44/51] Clarify mlflow docs

---
 docs/tutorials/qml/ml_tools.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/tutorials/qml/ml_tools.md b/docs/tutorials/qml/ml_tools.md
index c53f09e5c..727f5e1a0 100644
--- a/docs/tutorials/qml/ml_tools.md
+++ b/docs/tutorials/qml/ml_tools.md
@@ -323,7 +323,7 @@ export MLFLOW_EXPERIMENT=test_experiment
 export MLFLOW_RUN_NAME=run_0
 ```
 
-If no tracking URI is provided, mlflow stores runs and artifacts in the local `./mlflow` directory and if no names are defined, the experiment and run will be named with random UUIDs.
+If no tracking URI is provided, mlflow stores run information and artifacts in the local `./mlflow` directory and if no names are defined, the experiment and run will be named with random UUIDs.
 
 ### Setup
 Let's do the necessary imports and declare a `DataLoader`. We can already define some hyperparameters here, including the seed for random number generators. mlflow can log hyperparameters with arbitrary types, for example the observable that we want to monitor (`Z` in this case, which has a `qadence.Operation` type).

From d2a917b7212e92a4052269919b8c8e0e5f533715 Mon Sep 17 00:00:00 2001
From: Giorgio Tosti Balducci
 <64412197+debrevitatevitae@users.noreply.github.com>
Date: Wed, 17 Jul 2024 22:23:46 +0200
Subject: [PATCH 45/51] Lexical improvement in ml_tools.md

Co-authored-by: Roland-djee <9250798+Roland-djee@users.noreply.github.com>
---
 docs/tutorials/qml/ml_tools.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/tutorials/qml/ml_tools.md b/docs/tutorials/qml/ml_tools.md
index 727f5e1a0..f2d762ab9 100644
--- a/docs/tutorials/qml/ml_tools.md
+++ b/docs/tutorials/qml/ml_tools.md
@@ -366,7 +366,7 @@ def dataloader(batch_size: int = 25) -> DataLoader:
     return to_dataloader(x, y, batch_size=batch_size, infinite=True)
 ```
 
-We continue with the usual QNN definition, as well as loss function and optimizer.
+We continue with the regular QNN definition, together with the loss function and optimizer.
 
 ```python
 obs = hamiltonian_factory(register=hyperparams["n_qubits"], detuning=hyperparams["observable"])

From fee46267556502ab067c2c48e26ff7347b83f9d6 Mon Sep 17 00:00:00 2001
From: Giorgio Tosti Balducci
 <64412197+debrevitatevitae@users.noreply.github.com>
Date: Wed, 17 Jul 2024 22:24:59 +0200
Subject: [PATCH 46/51] Fix typing of InputData type

Co-authored-by: Roland-djee <9250798+Roland-djee@users.noreply.github.com>
---
 qadence/ml_tools/printing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index f9fcabadf..7d301d042 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -16,7 +16,7 @@
 logger = getLogger(__name__)
 
 PlottingFunction = Callable[[Module, int], tuple[str, Figure]]
-InputData = Union[Tensor, dict[str, Tensor]]
+InputData = Tensor | dict[str, Tensor]
 
 
 def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:

From aa98dfc220407fc59e463563f617748523d2d77d Mon Sep 17 00:00:00 2001
From: Giorgio Tosti Balducci
 <64412197+debrevitatevitae@users.noreply.github.com>
Date: Wed, 17 Jul 2024 22:26:13 +0200
Subject: [PATCH 47/51] Fix typing in log_model_mlflow

Co-authored-by: Roland-djee <9250798+Roland-djee@users.noreply.github.com>
---
 qadence/ml_tools/printing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 7d301d042..268ce7796 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -82,7 +82,7 @@ def plot_mlflow(
 def log_model_mlflow(
     writer: Any,
     model: Module,
-    dataloader: Union[None, DataLoader, DictDataLoader],
+    dataloader: DataLoader | DictDataLoader | None
 ) -> None:
     if dataloader is not None:
         xs: InputData

From 7df3c356ea82b9f7e84c5ea8dcfd70ebad085d27 Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 22:47:55 +0200
Subject: [PATCH 48/51] Fix typing with tracker mappings

---
 qadence/ml_tools/printing.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 7d301d042..9ae63677b 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -16,7 +16,7 @@
 logger = getLogger(__name__)
 
 PlottingFunction = Callable[[Module, int], tuple[str, Figure]]
-InputData = Tensor | dict[str, Tensor]
+InputData = Union[Tensor, dict[str, Tensor]]
 
 
 def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:
@@ -102,22 +102,22 @@ def log_model_mlflow(
     writer.pytorch.log_model(model, artifact_path="model", signature=signature)
 
 
-TRACKER_MAPPING = {
+TRACKER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
     ExperimentTrackingTool.TENSORBOARD: write_tensorboard,
     ExperimentTrackingTool.MLFLOW: write_mlflow,
 }
 
-LOGGER_MAPPING = {
+LOGGER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
     ExperimentTrackingTool.TENSORBOARD: log_hyperparams_tensorboard,
     ExperimentTrackingTool.MLFLOW: log_hyperparams_mlflow,
 }
 
-PLOTTER_MAPPING = {
+PLOTTER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
     ExperimentTrackingTool.TENSORBOARD: plot_tensorboard,
     ExperimentTrackingTool.MLFLOW: plot_mlflow,
 }
 
-MODEL_LOGGER_MAPPING = {
+MODEL_LOGGER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
     ExperimentTrackingTool.TENSORBOARD: log_model_tensorboard,
     ExperimentTrackingTool.MLFLOW: log_model_mlflow,
 }
@@ -126,7 +126,7 @@ def log_model_mlflow(
 def write_tracker(
     args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
 ) -> None:
-    return TRACKER_MAPPING[tracking_tool](*args)  # type: ignore
+    return TRACKER_MAPPING[tracking_tool](*args)
 
 
 def log_tracker(
@@ -138,10 +138,10 @@ def log_tracker(
 def plot_tracker(
     args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
 ) -> None:
-    return PLOTTER_MAPPING[tracking_tool](*args)  # type: ignore
+    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)  # type: ignore
+    return MODEL_LOGGER_MAPPING[tracking_tool](*args)

From 249dccb48c051e0a4deb9d7deee66cbb15f0d27d Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 22:54:56 +0200
Subject: [PATCH 49/51] Minor reformatting

---
 qadence/ml_tools/printing.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 2983312d8..70d661c9e 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -80,9 +80,7 @@ def plot_mlflow(
 
 
 def log_model_mlflow(
-    writer: Any,
-    model: Module,
-    dataloader: DataLoader | DictDataLoader | None
+    writer: Any, model: Module, dataloader: DataLoader | DictDataLoader | None
 ) -> None:
     if dataloader is not None:
         xs: InputData

From eb9a98cc9134fd95aafd7f48ec2e9ba1e161191b Mon Sep 17 00:00:00 2001
From: debrevitatevitae <giotb92@gmail.com>
Date: Wed, 17 Jul 2024 23:25:59 +0200
Subject: [PATCH 50/51] Turn args of trackers into variadic

---
 qadence/ml_tools/printing.py      |  8 ++++----
 qadence/ml_tools/train_grad.py    | 30 ++++++++++++++++++++----------
 qadence/ml_tools/train_no_grad.py | 15 ++++++++++-----
 3 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py
index 70d661c9e..047a0a8f8 100644
--- a/qadence/ml_tools/printing.py
+++ b/qadence/ml_tools/printing.py
@@ -122,24 +122,24 @@ def log_model_mlflow(
 
 
 def write_tracker(
-    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
 ) -> None:
     return TRACKER_MAPPING[tracking_tool](*args)
 
 
 def log_tracker(
-    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
 ) -> None:
     return LOGGER_MAPPING[tracking_tool](*args)
 
 
 def plot_tracker(
-    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+    *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
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
 ) -> None:
     return MODEL_LOGGER_MAPPING[tracking_tool](*args)
diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index a586baeaa..c607a8c1f 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -182,9 +182,7 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                 best_val_loss, metrics = loss_fn(model, xs_to_device)
 
                 metrics["val_loss"] = best_val_loss
-                write_tracker(
-                    (writer, None, metrics, init_iter), tracking_tool=config.tracking_tool
-                )
+                write_tracker(writer, None, metrics, init_iter, tracking_tool=config.tracking_tool)
 
             if config.folder:
                 if config.checkpoint_best_only:
@@ -193,7 +191,11 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                     write_checkpoint(config.folder, model, optimizer, init_iter)
 
             plot_tracker(
-                (writer, model, init_iter, config.plotting_functions), config.tracking_tool
+                writer,
+                model,
+                init_iter,
+                config.plotting_functions,
+                tracking_tool=config.tracking_tool,
             )
 
         except KeyboardInterrupt:
@@ -240,11 +242,17 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                     print_metrics(loss, metrics, iteration - 1)
 
                 if iteration % config.write_every == 0:
-                    write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+                    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), config.tracking_tool
+                        writer,
+                        model,
+                        iteration,
+                        config.plotting_functions,
+                        tracking_tool=config.tracking_tool,
                     )
                 if perform_val:
                     if iteration % config.val_every == 0:
@@ -256,7 +264,9 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
                             if config.folder and config.checkpoint_best_only:
                                 write_checkpoint(config.folder, model, optimizer, iteration="best")
                             metrics["val_loss"] = val_loss
-                            write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+                            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:
@@ -281,15 +291,15 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     # Final checkpointing and writing
     if config.folder and not config.checkpoint_best_only:
         write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+    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)
+        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)
+        log_model_tracker(writer, model, dataloader, tracking_tool=config.tracking_tool)
 
     # close tracker
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
diff --git a/qadence/ml_tools/train_no_grad.py b/qadence/ml_tools/train_no_grad.py
index 02fea1119..79684371a 100644
--- a/qadence/ml_tools/train_no_grad.py
+++ b/qadence/ml_tools/train_no_grad.py
@@ -56,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
@@ -121,11 +122,15 @@ def _update_parameters(
                 print_metrics(loss, metrics, iteration)
 
             if iteration % config.write_every == 0:
-                write_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+                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), config.tracking_tool
+                    writer,
+                    model,
+                    iteration,
+                    config.plotting_functions,
+                    tracking_tool=config.tracking_tool,
                 )
 
             if config.folder:
@@ -137,15 +142,15 @@ def _update_parameters(
 
     # writing hyperparameters
     if config.hyperparams:
-        log_tracker((writer, config.hyperparams, metrics), tracking_tool=config.tracking_tool)
+        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)
+        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_tracker((writer, loss, metrics, iteration), config.tracking_tool)
+    write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
 
     # close tracker
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:

From aa9d9a2d46258c8a1985b5d3b2433f39fed9fbe8 Mon Sep 17 00:00:00 2001
From: smit <smit.chaudhary@pasqal.com>
Date: Fri, 19 Jul 2024 12:31:32 +0200
Subject: [PATCH 51/51] remove stray comments

---
 qadence/ml_tools/train_grad.py    | 4 ----
 qadence/ml_tools/train_no_grad.py | 4 ----
 2 files changed, 8 deletions(-)

diff --git a/qadence/ml_tools/train_grad.py b/qadence/ml_tools/train_grad.py
index 5f512386e..89ac6383f 100644
--- a/qadence/ml_tools/train_grad.py
+++ b/qadence/ml_tools/train_grad.py
@@ -137,10 +137,6 @@ def loss_fn(model: torch.nn.Module, data: torch.Tensor) -> tuple[torch.Tensor, d
     else:
         writer = importlib.import_module("mlflow")
 
-        # writer.mlflow.pytorch.autolog(
-        #     log_every_n_step=config.write_every, log_models=False, log_datasets=False
-        # )
-
     perform_val = isinstance(config.val_every, int)
     if perform_val:
         if not isinstance(dataloader, DictDataLoader):
diff --git a/qadence/ml_tools/train_no_grad.py b/qadence/ml_tools/train_no_grad.py
index 79684371a..43a62dfe2 100644
--- a/qadence/ml_tools/train_no_grad.py
+++ b/qadence/ml_tools/train_no_grad.py
@@ -84,10 +84,6 @@ def _update_parameters(
     else:
         writer = importlib.import_module("mlflow")
 
-        # writer.mlflow.pytorch.autolog(
-        #     log_every_n_step=config.write_every, log_models=False, log_datasets=False
-        # )
-
     # set optimizer configuration and initial parameters
     optimizer.budget = config.max_iter
     optimizer.enable_pickling()