qadence 1.9.0__tar.gz → 1.9.2__tar.gz

{qadence-1.9.0 → qadence-1.9.2}/.github/workflows/test_fast.yml

@@ -36,12 +36,10 @@ jobs:
     - name: Run fast tests
       run: |
         hatch -v run test -m "not slow"
-    - name: Upload coverage data
-      uses: actions/upload-artifact@v4
+    - name: Upload coverage reports to Codecov
+      uses: codecov/codecov-action@v5
       with:
-        name: "coverage-data"
-        path: .coverage.*
-        if-no-files-found: ignore
+        token: ${{ secrets.CODECOV_TOKEN }}
 
   publish:
     name: Publish to PyPI

{qadence-1.9.0 → qadence-1.9.2}/PKG-INFO

@@ -1,9 +1,10 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: qadence
-Version: 1.9.0
+Version: 1.9.2
 Summary: Pasqal interface for circuit-based quantum computing SDKs
 Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>, Giorgio Tosti Balducci <giorgio.tosti-balducci@pasqal.com>, Daniele Cucurachi <daniele.cucurachi@pasqal.com>
 License: Apache 2.0
+License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
@@ -21,7 +22,7 @@ Requires-Dist: matplotlib
 Requires-Dist: nevergrad
 Requires-Dist: numpy
 Requires-Dist: openfermion
-Requires-Dist: pyqtorch==1.6.0
+Requires-Dist: pyqtorch==1.7.0
 Requires-Dist: pyyaml
 Requires-Dist: rich
 Requires-Dist: scipy
@@ -53,9 +54,9 @@ Requires-Dist: mlflow; extra == 'mlflow'
 Provides-Extra: protocols
 Requires-Dist: qadence-protocols; extra == 'protocols'
 Provides-Extra: pulser
-Requires-Dist: pasqal-cloud==0.12.5; extra == 'pulser'
-Requires-Dist: pulser-core==1.1.1; extra == 'pulser'
-Requires-Dist: pulser-simulation==1.1.1; extra == 'pulser'
+Requires-Dist: pasqal-cloud==0.12.6; extra == 'pulser'
+Requires-Dist: pulser-core==1.2.0; extra == 'pulser'
+Requires-Dist: pulser-simulation==1.2.0; extra == 'pulser'
 Provides-Extra: visualization
 Requires-Dist: graphviz; extra == 'visualization'
 Description-Content-Type: text/markdown
@@ -80,6 +81,8 @@ programs** with tunable qubit interactions and arbitrary register topologies rea
 [![Documentation](https://github.com/pasqal-io/qadence/actions/workflows/build_docs.yml/badge.svg)](https://pasqal-io.github.io/qadence/latest)
 [![Pypi](https://badge.fury.io/py/qadence.svg)](https://pypi.org/project/qadence/)
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+![Coverage](https://img.shields.io/codecov/c/github/pasqal-io/qadence?style=flat-square)
+
 
 ## Feature highlights
 

{qadence-1.9.0 → qadence-1.9.2}/README.md

@@ -18,6 +18,8 @@ programs** with tunable qubit interactions and arbitrary register topologies rea
 [![Documentation](https://github.com/pasqal-io/qadence/actions/workflows/build_docs.yml/badge.svg)](https://pasqal-io.github.io/qadence/latest)
 [![Pypi](https://badge.fury.io/py/qadence.svg)](https://pypi.org/project/qadence/)
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+![Coverage](https://img.shields.io/codecov/c/github/pasqal-io/qadence?style=flat-square)
+
 
 ## Feature highlights
 

{qadence-1.9.0 → qadence-1.9.2}/pyproject.toml

@@ -26,7 +26,7 @@ authors = [
 ]
 requires-python = ">=3.9"
 license = { text = "Apache 2.0" }
-version = "1.9.0"
+version = "1.9.2"
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python",
@@ -50,7 +50,7 @@ dependencies = [
     "jsonschema",
     "nevergrad",
     "scipy",
-    "pyqtorch==1.6.0",
+    "pyqtorch==1.7.0",
     "pyyaml",
     "matplotlib",
     "Arpeggio==2.0.2",
@@ -62,9 +62,9 @@ allow-ambiguous-features = true
 
 [project.optional-dependencies]
 pulser = [
-    "pulser-core==1.1.1",
-    "pulser-simulation==1.1.1",
-    "pasqal-cloud==0.12.5",
+    "pulser-core==1.2.0",
+    "pulser-simulation==1.2.0",
+    "pasqal-cloud==0.12.6",
 ]
 visualization = [
     "graphviz",
@@ -107,7 +107,7 @@ dependencies = [
 features = ["pulser", "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}"
+test = "pytest -n auto --cov-report=xml --cov-config=pyproject.toml --cov=qadence --cov=tests --ignore=./tests/test_examples.py {args}"
 test-examples = "pytest ./tests/test_examples.py {args}"
 no-cov = "cov --no-cov {args}"
 test-docs = "mkdocs build --clean --strict"
@@ -175,7 +175,7 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
 [tool.ruff]
 select = ["E", "F", "I", "Q"]
-extend-ignore = ["F841", "F403"]
+extend-ignore = ["F841", "F403", "E731", "E741"]
 line-length = 100
 
 [tool.ruff.isort]

{qadence-1.9.0 → qadence-1.9.2}/qadence/analog/device.py

@@ -5,6 +5,8 @@ from dataclasses import dataclass, fields
 from qadence.analog import AddressingPattern
 from qadence.types import PI, DeviceType, Interaction
 
+from .constants import C6_DICT
+
 
 @dataclass(frozen=True, eq=True)
 class RydbergDevice:
@@ -41,6 +43,11 @@ class RydbergDevice:
     type: DeviceType = DeviceType.IDEALIZED
     """DeviceType.IDEALIZED or REALISTIC to convert to the Pulser backend."""
 
+    @property
+    def coeff_ising(self) -> float:
+        """Value of C_6."""
+        return C6_DICT[self.rydberg_level]
+
     def __post_init__(self) -> None:
         # FIXME: Currently not supporting custom interaction functions.
         if self.interaction not in [Interaction.NN, Interaction.XY]:

{qadence-1.9.0 → qadence-1.9.2}/qadence/backends/pulser/backend.py

@@ -259,7 +259,7 @@ class Backend(BackendInterface):
             for i, param_values_el in enumerate(vals):
                 sequence = self.assign_parameters(circuit, param_values_el)
                 sim_result: CoherentResults = simulate_sequence(sequence, self.config, state)
-                final_state = sim_result.get_final_state().data.toarray()
+                final_state = sim_result.get_final_state().data.to_array()
                 batched_dm[i] = np.flip(final_state)
             return torch.from_numpy(batched_dm)
 

{qadence-1.9.0 → qadence-1.9.2}/qadence/backends/pyqtorch/convert_ops.py

@@ -264,7 +264,7 @@ def convert_block(
                 duration=duration,
                 solver=config.ode_solver,
                 steps=config.n_steps_hevo,
-                noise_operators=noise_operators,
+                noise=noise_operators if len(noise_operators) > 0 else None,
             )
         ]
 
@@ -351,22 +351,22 @@ def convert_block(
         )
 
 
-def convert_digital_noise(noise: NoiseHandler) -> pyq.noise.NoiseProtocol | None:
+def convert_digital_noise(noise: NoiseHandler) -> pyq.noise.DigitalNoiseProtocol | None:
     """Convert the digital noise into pyqtorch NoiseProtocol.
 
     Args:
         noise (NoiseHandler): Noise to convert.
 
     Returns:
-        pyq.noise.NoiseProtocol | None: Pyqtorch native noise protocol
+        pyq.noise.DigitalNoiseProtocol | None: Pyqtorch native noise protocol
             if there are any digital noise protocols.
     """
     digital_part = noise.filter(NoiseProtocol.DIGITAL)
     if digital_part is None:
         return None
-    return pyq.noise.NoiseProtocol(
+    return pyq.noise.DigitalNoiseProtocol(
         [
-            pyq.noise.NoiseProtocol(proto, option.get("error_probability"))
+            pyq.noise.DigitalNoiseProtocol(proto, option.get("error_probability"))
             for proto, option in zip(digital_part.protocol, digital_part.options)
         ]
     )

{qadence-1.9.0 → qadence-1.9.2}/qadence/backends/utils.py

@@ -110,9 +110,23 @@ def to_list_of_dicts(param_values: ParamDictType) -> list[ParamDictType]:
 
 
 def pyqify(state: Tensor, n_qubits: int = None) -> ArrayLike:
-    """Convert a state of shape (batch_size, 2**n_qubits) to [2] * n_qubits + [batch_size]."""
+    """Convert a state of shape (batch_size, 2**n_qubits) to [2] * n_qubits + [batch_size].
+
+    Or set the batch_size of a density matrix as the last dimension for PyQTorch.
+    """
     if n_qubits is None:
         n_qubits = int(log2(state.shape[1]))
+    if isinstance(state, DensityMatrix):
+        if (
+            len(state.shape) != 3
+            or (state.shape[1] != 2**n_qubits)
+            or (state.shape[1] != state.shape[2])
+        ):
+            raise ValueError(
+                "The initial state must be composed of tensors/arrays of size "
+                f"(batch_size, 2**n_qubits, 2**n_qubits). Found: {state.shape = }."
+            )
+        return torch.einsum("kij->ijk", state)
     if len(state.shape) != 2 or (state.shape[1] != 2**n_qubits):
         raise ValueError(
             "The initial state must be composed of tensors/arrays of size "

{qadence-1.9.0 → qadence-1.9.2}/qadence/engines/torch/differentiable_expectation.py

@@ -49,7 +49,9 @@ class PSRExpectation(Function):
             if isinstance(expectation_values[0], list):
                 exp_vals: list = []
                 for expectation_value in expectation_values:
-                    res = list(map(lambda x: x.get_final_state().data.toarray(), expectation_value))
+                    res = list(
+                        map(lambda x: x.get_final_state().data.to_array(), expectation_value)
+                    )
                     exp_vals.append(torch.tensor(res))
                 expectation_values = exp_vals
             return torch.stack(expectation_values)

{qadence-1.9.0 → qadence-1.9.2}/qadence/ml_tools/callbacks/__init__.py

@@ -2,9 +2,14 @@ from __future__ import annotations
 
 from .callback import (
     Callback,
+    EarlyStopping,
+    GradientMonitoring,
     LoadCheckpoint,
     LogHyperparameters,
     LogModelTracker,
+    LRSchedulerCosineAnnealing,
+    LRSchedulerCyclic,
+    LRSchedulerStepDecay,
     PlotMetrics,
     PrintMetrics,
     SaveBestCheckpoint,
@@ -26,5 +31,10 @@ __all__ = [
     "SaveBestCheckpoint",
     "SaveCheckpoint",
     "WriteMetrics",
+    "GradientMonitoring",
+    "LRSchedulerStepDecay",
+    "LRSchedulerCyclic",
+    "LRSchedulerCosineAnnealing",
+    "EarlyStopping",
     "get_writer",
 ]

{qadence-1.9.0 → qadence-1.9.2}/qadence/ml_tools/callbacks/callback.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import math
+from logging import getLogger
 from typing import Any, Callable
 
 from qadence.ml_tools.callbacks.saveload import load_checkpoint, write_checkpoint
@@ -12,6 +14,8 @@ from qadence.ml_tools.stages import TrainingStage
 CallbackFunction = Callable[..., Any]
 CallbackConditionFunction = Callable[..., bool]
 
+logger = getLogger("ml_tools")
+
 
 class Callback:
     """Base class for defining various training callbacks.
@@ -258,7 +262,7 @@ class WriteMetrics(Callback):
             writer (BaseWriter ): The writer object for logging.
         """
         opt_result = trainer.opt_result
-        writer.write(opt_result)
+        writer.write(opt_result.iteration, opt_result.metrics)
 
 
 class PlotMetrics(Callback):
@@ -449,3 +453,323 @@ class LogModelTracker(Callback):
         writer.log_model(
             model, trainer.train_dataloader, trainer.val_dataloader, trainer.test_dataloader
         )
+
+
+class LRSchedulerStepDecay(Callback):
+    """
+    Reduces the learning rate by a factor at regular intervals.
+
+    This callback adjusts the learning rate by multiplying it with a decay factor
+    after a specified number of iterations. The learning rate is updated as:
+        lr = lr * gamma
+
+    Example Usage in `TrainConfig`:
+    To use `LRSchedulerStepDecay`, include it in the `callbacks` list when setting
+    up your `TrainConfig`:
+    ```python exec="on" source="material-block" result="json"
+    from qadence.ml_tools import TrainConfig
+    from qadence.ml_tools.callbacks import LRSchedulerStepDecay
+
+    # Create an instance of the LRSchedulerStepDecay callback
+    lr_step_decay = LRSchedulerStepDecay(on="train_epoch_end",
+                                         called_every=100,
+                                         gamma=0.5)
+
+    config = TrainConfig(
+        max_iter=10000,
+        # Print metrics every 1000 training epochs
+        print_every=1000,
+        # Add the custom callback
+        callbacks=[lr_step_decay]
+    )
+    ```
+    """
+
+    def __init__(self, on: str, called_every: int, gamma: float = 0.5):
+        """Initializes the LRSchedulerStepDecay callback.
+
+        Args:
+            on (str): The event to trigger the callback.
+            called_every (int): Frequency of callback calls in terms of iterations.
+            gamma (float, optional): The decay factor applied to the learning rate.
+                A value < 1 reduces the learning rate over time. Default is 0.5.
+        """
+        super().__init__(on=on, called_every=called_every)
+        self.gamma = gamma
+
+    def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
+        """
+        Runs the callback to apply step decay to the learning rate.
+
+        Args:
+            trainer (Any): The training object.
+            config (TrainConfig): The configuration object.
+            writer (BaseWriter): The writer object for logging.
+        """
+        for param_group in trainer.optimizer.param_groups:
+            param_group["lr"] *= self.gamma
+
+
+class LRSchedulerCyclic(Callback):
+    """
+    Applies a cyclic learning rate schedule during training.
+
+    This callback oscillates the learning rate between a minimum (base_lr)
+    and a maximum (max_lr) over a defined cycle length (step_size). The learning
+    rate follows a triangular wave pattern.
+
+    Example Usage in `TrainConfig`:
+    To use `LRSchedulerCyclic`, include it in the `callbacks` list when setting
+    up your `TrainConfig`:
+    ```python exec="on" source="material-block" result="json"
+    from qadence.ml_tools import TrainConfig
+    from qadence.ml_tools.callbacks import LRSchedulerCyclic
+
+    # Create an instance of the LRSchedulerCyclic callback
+    lr_cyclic = LRSchedulerCyclic(on="train_batch_end",
+                                  called_every=1,
+                                  base_lr=0.001,
+                                  max_lr=0.01,
+                                  step_size=2000)
+
+    config = TrainConfig(
+        max_iter=10000,
+        # Print metrics every 1000 training epochs
+        print_every=1000,
+        # Add the custom callback
+        callbacks=[lr_cyclic]
+    )
+    ```
+    """
+
+    def __init__(self, on: str, called_every: int, base_lr: float, max_lr: float, step_size: int):
+        """Initializes the LRSchedulerCyclic callback.
+
+        Args:
+            on (str): The event to trigger the callback.
+            called_every (int): Frequency of callback calls in terms of iterations.
+            base_lr (float): The minimum learning rate.
+            max_lr (float): The maximum learning rate.
+            step_size (int): Number of iterations for half a cycle.
+        """
+        super().__init__(on=on, called_every=called_every)
+        self.base_lr = base_lr
+        self.max_lr = max_lr
+        self.step_size = step_size
+
+    def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
+        """
+        Adjusts the learning rate cyclically.
+
+        Args:
+            trainer (Any): The training object.
+            config (TrainConfig): The configuration object.
+            writer (BaseWriter): The writer object for logging.
+        """
+        cycle = trainer.opt_result.iteration // (2 * self.step_size)
+        x = abs(trainer.opt_result.iteration / self.step_size - 2 * cycle - 1)
+        scale = max(0, (1 - x))
+        new_lr = self.base_lr + (self.max_lr - self.base_lr) * scale
+        for param_group in trainer.optimizer.param_groups:
+            param_group["lr"] = new_lr
+
+
+class LRSchedulerCosineAnnealing(Callback):
+    """
+    Applies cosine annealing to the learning rate during training.
+
+    This callback decreases the learning rate following a cosine curve,
+    starting from the initial learning rate and annealing to a minimum (min_lr).
+
+    Example Usage in `TrainConfig`:
+    To use `LRSchedulerCosineAnnealing`, include it in the `callbacks` list
+    when setting up your `TrainConfig`:
+    ```python exec="on" source="material-block" result="json"
+    from qadence.ml_tools import TrainConfig
+    from qadence.ml_tools.callbacks import LRSchedulerCosineAnnealing
+
+    # Create an instance of the LRSchedulerCosineAnnealing callback
+    lr_cosine = LRSchedulerCosineAnnealing(on="train_batch_end",
+                                           called_every=1,
+                                           t_max=5000,
+                                           min_lr=1e-6)
+
+    config = TrainConfig(
+        max_iter=10000,
+        # Print metrics every 1000 training epochs
+        print_every=1000,
+        # Add the custom callback
+        callbacks=[lr_cosine]
+    )
+    ```
+    """
+
+    def __init__(self, on: str, called_every: int, t_max: int, min_lr: float = 0.0):
+        """Initializes the LRSchedulerCosineAnnealing callback.
+
+        Args:
+            on (str): The event to trigger the callback.
+            called_every (int): Frequency of callback calls in terms of iterations.
+            t_max (int): The total number of iterations for one annealing cycle.
+            min_lr (float, optional): The minimum learning rate. Default is 0.0.
+        """
+        super().__init__(on=on, called_every=called_every)
+        self.t_max = t_max
+        self.min_lr = min_lr
+
+    def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
+        """
+        Adjusts the learning rate using cosine annealing.
+
+        Args:
+            trainer (Any): The training object.
+            config (TrainConfig): The configuration object.
+            writer (BaseWriter): The writer object for logging.
+        """
+        for param_group in trainer.optimizer.param_groups:
+            max_lr = param_group["lr"]
+            new_lr = (
+                self.min_lr
+                + (max_lr - self.min_lr)
+                * (1 + math.cos(math.pi * trainer.opt_result.iteration / self.t_max))
+                / 2
+            )
+            param_group["lr"] = new_lr
+
+
+class EarlyStopping(Callback):
+    """
+    Stops training when a monitored metric has not improved for a specified number of epochs.
+
+    This callback monitors a specified metric (e.g., validation loss or accuracy). If the metric
+    does not improve for a given patience period, training is stopped.
+
+    Example Usage in `TrainConfig`:
+    To use `EarlyStopping`, include it in the `callbacks` list when setting up your `TrainConfig`:
+    ```python exec="on" source="material-block" result="json"
+    from qadence.ml_tools import TrainConfig
+    from qadence.ml_tools.callbacks import EarlyStopping
+
+    # Create an instance of the EarlyStopping callback
+    early_stopping = EarlyStopping(on="val_epoch_end",
+                                   called_every=1,
+                                   monitor="val_loss",
+                                   patience=5,
+                                   mode="min")
+
+    config = TrainConfig(
+        max_iter=10000,
+        print_every=1000,
+        callbacks=[early_stopping]
+    )
+    ```
+    """
+
+    def __init__(
+        self, on: str, called_every: int, monitor: str, patience: int = 5, mode: str = "min"
+    ):
+        """Initializes the EarlyStopping callback.
+
+        Args:
+            on (str): The event to trigger the callback (e.g., "val_epoch_end").
+            called_every (int): Frequency of callback calls in terms of iterations.
+            monitor (str): The metric to monitor (e.g., "val_loss" or "train_loss").
+                All metrics returned by optimize step are available to monitor.
+                Please add "val_" and "train_" strings at the start of the metric name.
+            patience (int, optional): Number of iterations to wait for improvement. Default is 5.
+            mode (str, optional): Whether to minimize ("min") or maximize ("max") the metric.
+                Default is "min".
+        """
+        super().__init__(on=on, called_every=called_every)
+        self.monitor = monitor
+        self.patience = patience
+        self.mode = mode
+        self.best_value = float("inf") if mode == "min" else -float("inf")
+        self.counter = 0
+
+    def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
+        """
+        Monitors the metric and stops training if no improvement is observed.
+
+        Args:
+            trainer (Any): The training object.
+            config (TrainConfig): The configuration object.
+            writer (BaseWriter): The writer object for logging.
+        """
+        current_value = trainer.opt_result.metrics.get(self.monitor)
+        if current_value is None:
+            raise ValueError(f"Metric '{self.monitor}' is not available in the trainer's metrics.")
+
+        if (self.mode == "min" and current_value < self.best_value) or (
+            self.mode == "max" and current_value > self.best_value
+        ):
+            self.best_value = current_value
+            self.counter = 0
+        else:
+            self.counter += 1
+
+        if self.counter >= self.patience:
+            logger.info(
+                f"EarlyStopping: No improvement in '{self.monitor}' for {self.patience} epochs. "
+                "Stopping training."
+            )
+            trainer.stop_training = True
+
+
+class GradientMonitoring(Callback):
+    """
+    Logs gradient statistics (e.g., mean, standard deviation, max) during training.
+
+    This callback monitors and logs statistics about the gradients of the model parameters
+    to help debug or optimize the training process.
+
+    Example Usage in `TrainConfig`:
+    To use `GradientMonitoring`, include it in the `callbacks` list when
+    setting up your `TrainConfig`:
+    ```python exec="on" source="material-block" result="json"
+    from qadence.ml_tools import TrainConfig
+    from qadence.ml_tools.callbacks import GradientMonitoring
+
+    # Create an instance of the GradientMonitoring callback
+    gradient_monitoring = GradientMonitoring(on="train_batch_end", called_every=10)
+
+    config = TrainConfig(
+        max_iter=10000,
+        print_every=1000,
+        callbacks=[gradient_monitoring]
+    )
+    ```
+    """
+
+    def __init__(self, on: str, called_every: int = 1):
+        """Initializes the GradientMonitoring callback.
+
+        Args:
+            on (str): The event to trigger the callback (e.g., "train_batch_end").
+            called_every (int): Frequency of callback calls in terms of iterations.
+        """
+        super().__init__(on=on, called_every=called_every)
+
+    def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> None:
+        """
+        Logs gradient statistics.
+
+        Args:
+            trainer (Any): The training object.
+            config (TrainConfig): The configuration object.
+            writer (BaseWriter): The writer object for logging.
+        """
+        gradient_stats = {}
+        for name, param in trainer.model.named_parameters():
+            if param.grad is not None:
+                grad = param.grad
+                gradient_stats.update(
+                    {
+                        name + "_mean": grad.mean().item(),
+                        name + "_std": grad.std().item(),
+                        name + "_max": grad.max().item(),
+                        name + "_min": grad.min().item(),
+                    }
+                )
+
+        writer.write(trainer.opt_result.iteration, gradient_stats)