qadence 1.10.3__py3-none-any.whl → 1.11.1__py3-none-any.whl

qadence/blocks/block_to_tensor.py

@@ -82,18 +82,20 @@ def _fill_identities(
     full_qubit_support = tuple(sorted(full_qubit_support))
     qubit_support = tuple(sorted(qubit_support))
     block_mat = block_mat.to(device)
-    mat = IMAT.to(device) if qubit_support[0] != full_qubit_support[0] else block_mat
+    identity_mat = IMAT.to(device)
     if diag_only:
-        mat = torch.diag(mat.squeeze(0))
+        block_mat = torch.diag(block_mat.squeeze(0))
+        identity_mat = torch.diag(identity_mat.squeeze(0))
+    mat = identity_mat if qubit_support[0] != full_qubit_support[0] else block_mat
     for i in full_qubit_support[1:]:
         if i == qubit_support[0]:
-            other = torch.diag(block_mat.squeeze(0)) if diag_only else block_mat
+            other = block_mat
             if endianness == Endianness.LITTLE:
                 mat = torch.kron(other, mat)
             else:
                 mat = torch.kron(mat.contiguous(), other.contiguous())
         elif i not in qubit_support:
-            other = torch.diag(IMAT.squeeze(0).to(device)) if diag_only else IMAT.to(device)
+            other = identity_mat
             if endianness == Endianness.LITTLE:
                 mat = torch.kron(other.contiguous(), mat.contiguous())
             else:
@@ -264,13 +266,12 @@ def _gate_parameters(b: AbstractBlock, values: dict[str, torch.Tensor]) -> tuple
 
 def block_to_diagonal(
     block: AbstractBlock,
+    values: dict[str, TNumber | torch.Tensor] = dict(),
     qubit_support: tuple | list | None = None,
-    use_full_support: bool = True,
+    use_full_support: bool = False,
     endianness: Endianness = Endianness.BIG,
     device: torch.device = None,
 ) -> torch.Tensor:
-    if block.is_parametric:
-        raise TypeError("Sparse observables cant be parametric.")
     if not block._is_diag_pauli:
         raise TypeError("Sparse observables can only be used on paulis which are diagonal.")
     if qubit_support is None:
@@ -282,17 +283,16 @@ def block_to_diagonal(
     if isinstance(block, (ChainBlock, KronBlock)):
         v = torch.ones(2**nqubits, dtype=torch.cdouble)
         for b in block.blocks:
-            v *= block_to_diagonal(b, qubit_support)
+            v *= block_to_diagonal(b, values, qubit_support, device=device)
     if isinstance(block, AddBlock):
         t = torch.zeros(2**nqubits, dtype=torch.cdouble)
         for b in block.blocks:
-            t += block_to_diagonal(b, qubit_support)
+            t += block_to_diagonal(b, values, qubit_support, device=device)
         v = t
     elif isinstance(block, ScaleBlock):
-        _s = evaluate(block.scale, {}, as_torch=True)  # type: ignore[attr-defined]
-        _s = _s.detach()  # type: ignore[union-attr]
-        v = _s * block_to_diagonal(block.block, qubit_support)
-
+        _s = evaluate(block.scale, values, as_torch=True)  # type: ignore[attr-defined]
+        _s = _s.detach().squeeze(0)  # type: ignore[union-attr]
+        v = _s * block_to_diagonal(block.block, values, qubit_support, device=device)
     elif isinstance(block, PrimitiveBlock):
         v = _fill_identities(
             OPERATIONS_DICT[block.name],
@@ -300,6 +300,7 @@ def block_to_diagonal(
             qubit_support,  # type: ignore [arg-type]
             diag_only=True,
             endianness=endianness,
+            device=device,
         )
     return v
 
@@ -309,7 +310,7 @@ def block_to_tensor(
     block: AbstractBlock,
     values: dict[str, TNumber | torch.Tensor] = {},
     qubit_support: tuple | None = None,
-    use_full_support: bool = True,
+    use_full_support: bool = False,
     tensor_type: TensorType = TensorType.DENSE,
     endianness: Endianness = Endianness.BIG,
     device: torch.device = None,
@@ -339,18 +340,14 @@ def block_to_tensor(
     print(block_to_tensor(obs, tensor_type="SparseDiagonal"))
     ```
     """
+    from qadence.blocks import embedding
 
-    # FIXME: default use_full_support to False. In general, it would
-    # be more efficient to do that, and make sure that computations such
-    # as observables only do the matmul of the size of the qubit support.
-
+    (ps, embed) = embedding(block)
+    values = embed(ps, values)
     if tensor_type == TensorType.DENSE:
-        from qadence.blocks import embedding
-
-        (ps, embed) = embedding(block)
         return _block_to_tensor_embedded(
             block,
-            embed(ps, values),
+            values,
             qubit_support,
             use_full_support,
             endianness=endianness,
@@ -358,7 +355,7 @@ def block_to_tensor(
         )
 
     elif tensor_type == TensorType.SPARSEDIAGONAL:
-        t = block_to_diagonal(block, endianness=endianness)
+        t = block_to_diagonal(block, values, endianness=endianness)
         indices, values, size = torch.nonzero(t), t[t != 0], len(t)
         indices = torch.stack((indices.flatten(), indices.flatten()))
         return torch.sparse_coo_tensor(indices, values, (size, size))
@@ -369,7 +366,7 @@ def _block_to_tensor_embedded(
     block: AbstractBlock,
     values: dict[str, TNumber | torch.Tensor] = {},
     qubit_support: tuple | None = None,
-    use_full_support: bool = True,
+    use_full_support: bool = False,
     endianness: Endianness = Endianness.BIG,
     device: torch.device = None,
 ) -> torch.Tensor:

qadence/constructors/__init__.py

@@ -17,6 +17,9 @@ from .hamiltonians import (
     ObservableConfig,
     total_magnetization,
     zz_hamiltonian,
+    total_magnetization_config,
+    zz_hamiltonian_config,
+    ising_hamiltonian_config,
 )
 
 from .rydberg_hea import rydberg_hea, rydberg_hea_layer
@@ -34,9 +37,12 @@ __all__ = [
     "iia",
     "hamiltonian_factory",
     "ising_hamiltonian",
-    "ObservableConfig",
     "total_magnetization",
     "zz_hamiltonian",
+    "ObservableConfig",
+    "total_magnetization_config",
+    "zz_hamiltonian_config",
+    "ising_hamiltonian_config",
     "qft",
     "daqc_transform",
     "rydberg_hea",

qadence/constructors/hamiltonians.py

@@ -7,11 +7,12 @@ from typing import Callable, List, Type, Union
 import numpy as np
 from torch import Tensor, double, ones, rand
 from typing_extensions import Any
+from qadence.parameters import Parameter
 
 from qadence.blocks import AbstractBlock, add, block_is_qubit_hamiltonian
-from qadence.operations import N, X, Y, Z
+from qadence.operations import N, X, Y, Z, H
 from qadence.register import Register
-from qadence.types import Interaction, ObservableTransform, TArray, TParameter
+from qadence.types import Interaction, TArray, TParameter
 
 logger = getLogger(__name__)
 
@@ -239,7 +240,30 @@ def is_numeric(x: Any) -> bool:
 
 @dataclass
 class ObservableConfig:
-    detuning: TDetuning
+    """ObservableConfig is a configuration class for defining the parameters of an observable Hamiltonian."""
+
+    interaction: Interaction | Callable | None = None
+    """
+    The type of interaction.
+
+    Available options from the Interaction enum are:
+            - Interaction.ZZ
+            - Interaction.NN
+            - Interaction.XY
+            - Interaction.XYZ
+
+    Alternatively, a custom interaction function can be defined.
+            Example:
+
+                def custom_int(i: int, j: int):
+                    return X(i) @ X(j) + Y(i) @ Y(j)
+
+                n_qubits = 2
+
+                observable_config = ObservableConfig(interaction=custom_int, scale = 1.0, shift = 0.0)
+                observable = create_observable(register=4, config=observable_config)
+    """
+    detuning: TDetuning | None = None
     """
     Single qubit detuning of the observable Hamiltonian.
 
@@ -249,8 +273,6 @@ class ObservableConfig:
     """The scale by which to multiply the output of the observable."""
     shift: TParameter = 0.0
     """The shift to add to the output of the observable."""
-    transformation_type: ObservableTransform = ObservableTransform.NONE  # type: ignore[assignment]
-    """The type of transformation."""
     trainable_transform: bool | None = None
     """
     Whether to have a trainable transformation on the output of the observable.
@@ -259,10 +281,84 @@ class ObservableConfig:
     If True, the scale and shift are VariationalParameter.
     If False, the scale and shift are FeatureParameter.
     """
+    tag: str | None = None
+    """
+    String to indicate the name tag of the observable.
+
+    Defaults to None, in which case no tag will be applied.
+    """
 
     def __post_init__(self) -> None:
+        if self.interaction is None and self.detuning is None:
+            raise ValueError(
+                "Please provide an interaction and/or detuning for the Observable Hamiltonian."
+            )
+
         if is_numeric(self.scale) and is_numeric(self.shift):
-            assert (
-                self.trainable_transform is None
-            ), f"If scale and shift are numbers, trainable_transform must be None. \
-            But got: {self.trainable_transform}"
+            assert self.trainable_transform is None, (
+                "If scale and shift are numbers, trainable_transform must be None."
+                f"But got: {self.trainable_transform}"
+            )
+
+        # trasform the scale and shift into parameters
+        if self.trainable_transform is not None:
+            self.shift = Parameter(name=self.shift, trainable=self.trainable_transform)
+            self.scale = Parameter(name=self.scale, trainable=self.trainable_transform)
+        else:
+            self.shift = Parameter(self.shift)
+            self.scale = Parameter(self.scale)
+
+
+def total_magnetization_config(
+    scale: TParameter = 1.0,
+    shift: TParameter = 0.0,
+    trainable_transform: bool | None = None,
+) -> ObservableConfig:
+    return ObservableConfig(
+        detuning=Z,
+        scale=scale,
+        shift=shift,
+        trainable_transform=trainable_transform,
+        tag="Total Magnetization",
+    )
+
+
+def zz_hamiltonian_config(
+    scale: TParameter = 1.0,
+    shift: TParameter = 0.0,
+    trainable_transform: bool | None = None,
+) -> ObservableConfig:
+    return ObservableConfig(
+        interaction=Interaction.ZZ,
+        detuning=Z,
+        scale=scale,
+        shift=shift,
+        trainable_transform=trainable_transform,
+        tag="ZZ Hamiltonian",
+    )
+
+
+def ising_hamiltonian_config(
+    scale: TParameter = 1.0,
+    shift: TParameter = 0.0,
+    trainable_transform: bool | None = None,
+) -> ObservableConfig:
+
+    def ZZ_Z_hamiltonian(i: int, j: int) -> AbstractBlock:
+        result = Z(i) @ Z(j)
+
+        if i == 0:
+            result += Z(j)
+        elif i == 1 and j == 2:
+            result += Z(0)
+
+        return result
+
+    return ObservableConfig(
+        interaction=ZZ_Z_hamiltonian,
+        detuning=X,
+        scale=scale,
+        shift=shift,
+        trainable_transform=trainable_transform,
+        tag="Ising Hamiltonian",
+    )

qadence/draw/utils.py

@@ -147,7 +147,8 @@ def _(model: QuantumModel, *args: Any, **kwargs: Any) -> QuantumCircuitDiagram:
             raise ValueError("Cannot visualize QuantumModel with more than one observable.")
 
         obs = deepcopy(model._observable[0].original)  # type: ignore [index]
-        obs.tag = "Obs."
+        if not isinstance(obs.tag, str):
+            obs.tag = "Obs."
 
         block: AbstractBlock = chain(model._circuit.original.block, obs)
     else:

qadence/mitigations/analog_zne.py

@@ -92,7 +92,9 @@ def pulse_experiment(
         )
     # Convert observable to Numpy types compatible with QuTip simulations.
     # Matrices are flipped to match QuTip conventions.
-    converted_observable = [np.flip(block_to_tensor(obs).numpy()) for obs in observable]
+    converted_observable = [
+        np.flip(block_to_tensor(obs, use_full_support=True).numpy()) for obs in observable
+    ]
     # Create ZNE datasets by looping over batches.
     for observable in converted_observable:
         # Get expectation values at the end of the time serie [0,t]
@@ -130,7 +132,9 @@ def noise_level_experiment(
     )
     # Convert observable to Numpy types compatible with QuTip simulations.
     # Matrices are flipped to match QuTip conventions.
-    converted_observable = [np.flip(block_to_tensor(obs).numpy()) for obs in observable]
+    converted_observable = [
+        np.flip(block_to_tensor(obs, use_full_support=True).numpy()) for obs in observable
+    ]
     # Create ZNE datasets by looping over batches.
     for observable in converted_observable:
         # Get expectation values at the end of the time serie [0,t]

qadence/ml_tools/__init__.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from .callbacks.saveload import load_checkpoint, load_model, write_checkpoint
 from .config import AnsatzConfig, FeatureMapConfig, TrainConfig
-from .constructors import create_ansatz, create_fm_blocks, observable_from_config
+from .constructors import create_ansatz, create_fm_blocks, create_observable
 from .data import DictDataLoader, InfiniteTensorDataset, OptimizeResult, to_dataloader
 from .information import InformationContent
 from .models import QNN
@@ -19,7 +19,7 @@ __all__ = [
     "DictDataLoader",
     "FeatureMapConfig",
     "load_checkpoint",
-    "observable_from_config",
+    "create_observable",
     "QNN",
     "TrainConfig",
     "OptimizeResult",

qadence/ml_tools/callbacks/callback.py

@@ -95,14 +95,36 @@ class Callback:
         self.callback: CallbackFunction | None = callback
         self.on: str | TrainingStage = on
         self.called_every: int = called_every
-        self.callback_condition = callback_condition or (lambda _: True)
+        self.callback_condition = (
+            callback_condition if callback_condition else Callback.default_callback
+        )
 
         if isinstance(modify_optimize_result, dict):
-            self.modify_optimize_result = (
-                lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
+            self.modify_optimize_result = lambda opt_res: Callback.modify_opt_res_dict(
+                opt_res, modify_optimize_result
             )
         else:
-            self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)
+            self.modify_optimize_result = (
+                modify_optimize_result
+                if modify_optimize_result
+                else Callback.modify_opt_res_default
+            )
+
+    @staticmethod
+    def default_callback(_: Any) -> bool:
+        return True
+
+    @staticmethod
+    def modify_opt_res_dict(
+        opt_res: OptimizeResult,
+        modify_optimize_result: dict[str, Any] = {},
+    ) -> OptimizeResult:
+        opt_res.extra.update(modify_optimize_result)
+        return opt_res
+
+    @staticmethod
+    def modify_opt_res_default(opt_res: OptimizeResult) -> OptimizeResult:
+        return opt_res
 
     @property
     def on(self) -> TrainingStage | str:
@@ -261,8 +283,9 @@ class WriteMetrics(Callback):
             config (TrainConfig): The configuration object.
             writer (BaseWriter ): The writer object for logging.
         """
-        opt_result = trainer.opt_result
-        writer.write(opt_result.iteration, opt_result.metrics)
+        if trainer.accelerator.rank == 0:
+            opt_result = trainer.opt_result
+            writer.write(opt_result.iteration, opt_result.metrics)
 
 
 class PlotMetrics(Callback):
@@ -299,9 +322,10 @@ class PlotMetrics(Callback):
             config (TrainConfig): The configuration object.
             writer (BaseWriter ): The writer object for logging.
         """
-        opt_result = trainer.opt_result
-        plotting_functions = config.plotting_functions
-        writer.plot(trainer.model, opt_result.iteration, plotting_functions)
+        if trainer.accelerator.rank == 0:
+            opt_result = trainer.opt_result
+            plotting_functions = config.plotting_functions
+            writer.plot(trainer.model, opt_result.iteration, plotting_functions)
 
 
 class LogHyperparameters(Callback):
@@ -338,8 +362,9 @@ class LogHyperparameters(Callback):
             config (TrainConfig): The configuration object.
             writer (BaseWriter ): The writer object for logging.
         """
-        hyperparams = config.hyperparams
-        writer.log_hyperparams(hyperparams)
+        if trainer.accelerator.rank == 0:
+            hyperparams = config.hyperparams
+            writer.log_hyperparams(hyperparams)
 
 
 class SaveCheckpoint(Callback):
@@ -376,11 +401,12 @@ class SaveCheckpoint(Callback):
             config (TrainConfig): The configuration object.
             writer (BaseWriter ): The writer object for logging.
         """
-        folder = config.log_folder
-        model = trainer.model
-        optimizer = trainer.optimizer
-        opt_result = trainer.opt_result
-        write_checkpoint(folder, model, optimizer, opt_result.iteration)
+        if trainer.accelerator.rank == 0:
+            folder = config.log_folder
+            model = trainer.model
+            optimizer = trainer.optimizer
+            opt_result = trainer.opt_result
+            write_checkpoint(folder, model, optimizer, opt_result.iteration)
 
 
 class SaveBestCheckpoint(SaveCheckpoint):
@@ -404,17 +430,18 @@ class SaveBestCheckpoint(SaveCheckpoint):
             config (TrainConfig): The configuration object.
             writer (BaseWriter ): The writer object for logging.
         """
-        opt_result = trainer.opt_result
-        if config.validation_criterion and config.validation_criterion(
-            opt_result.loss, self.best_loss, config.val_epsilon
-        ):
-            self.best_loss = opt_result.loss
-
-            folder = config.log_folder
-            model = trainer.model
-            optimizer = trainer.optimizer
+        if trainer.accelerator.rank == 0:
             opt_result = trainer.opt_result
-            write_checkpoint(folder, model, optimizer, "best")
+            if config.validation_criterion and config.validation_criterion(
+                opt_result.loss, self.best_loss, config.val_epsilon
+            ):
+                self.best_loss = opt_result.loss
+
+                folder = config.log_folder
+                model = trainer.model
+                optimizer = trainer.optimizer
+                opt_result = trainer.opt_result
+                write_checkpoint(folder, model, optimizer, "best")
 
 
 class LoadCheckpoint(Callback):
@@ -431,11 +458,12 @@ class LoadCheckpoint(Callback):
         Returns:
             Any: The result of loading the checkpoint.
         """
-        folder = config.log_folder
-        model = trainer.model
-        optimizer = trainer.optimizer
-        device = trainer.log_device
-        return load_checkpoint(folder, model, optimizer, device=device)
+        if trainer.accelerator.rank == 0:
+            folder = config.log_folder
+            model = trainer.model
+            optimizer = trainer.optimizer
+            device = trainer.accelerator.execution.log_device
+            return load_checkpoint(folder, model, optimizer, device=device)
 
 
 class LogModelTracker(Callback):
@@ -449,10 +477,11 @@ class LogModelTracker(Callback):
             config (TrainConfig): The configuration object.
             writer (BaseWriter ): The writer object for logging.
         """
-        model = trainer.model
-        writer.log_model(
-            model, trainer.train_dataloader, trainer.val_dataloader, trainer.test_dataloader
-        )
+        if trainer.accelerator.rank == 0:
+            model = trainer.model
+            writer.log_model(
+                model, trainer.train_dataloader, trainer.val_dataloader, trainer.test_dataloader
+            )
 
 
 class LRSchedulerStepDecay(Callback):
@@ -713,7 +742,7 @@ class EarlyStopping(Callback):
                 f"EarlyStopping: No improvement in '{self.monitor}' for {self.patience} epochs. "
                 "Stopping training."
             )
-            trainer.stop_training = True
+            trainer._stop_training.fill_(1)
 
 
 class GradientMonitoring(Callback):
@@ -759,17 +788,18 @@ class GradientMonitoring(Callback):
             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)
+        if trainer.accelerator.rank == 0:
+            gradient_stats = {}
+            for name, param in trainer.model.named_parameters():
+                if param.grad is not None:
+                    grad = param.grad
+                    gradient_stats.update(
+                        {
+                            name + "_mean": grad.mean().item(),
+                            name + "_std": grad.std().item(),
+                            name + "_max": grad.max().item(),
+                            name + "_min": grad.min().item(),
+                        }
+                    )
+
+            writer.write(trainer.opt_result.iteration, gradient_stats)

qadence/ml_tools/callbacks/callbackmanager.py

@@ -201,7 +201,8 @@ class CallbacksManager:
                     logger.debug(f"Loaded model and optimizer from {self.config.log_folder}")
 
         # Setup writer
-        self.writer.open(self.config, iteration=trainer.global_step)
+        if trainer.accelerator.rank == 0:
+            self.writer.open(self.config, iteration=trainer.global_step)
 
     def end_training(self, trainer: Any) -> None:
         """
@@ -210,5 +211,5 @@ class CallbacksManager:
         Args:
             trainer (Any): The training object managing the training process.
         """
-        if self.writer:
+        if trainer.accelerator.rank == 0 and self.writer:
             self.writer.close()

qadence/ml_tools/callbacks/writer_registry.py

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

qadence/ml_tools/config.py

@@ -20,6 +20,7 @@ from qadence.types import (
     ReuploadScaling,
     Strategy,
 )
+from torch import dtype
 
 logger = getLogger(__file__)
 
@@ -116,10 +117,9 @@ class TrainConfig:
     """The log folder for saving checkpoints and tensorboard logs.
 
     This stores the path where all logs and checkpoints are being saved
-    for this training session. `log_folder` takes precedence over `root_folder` and
-    `create_subfolder_per_run` arguments. If the user specifies a log_folder,
-    all checkpoints will be saved in this folder and `root_folder` argument
-    will not be used.
+    for this training session. `log_folder` takes precedence over `root_folder`,
+    but it is ignored if `create_subfolders_per_run=True` (in which case, subfolders
+    will be spawned in the root folder).
     """
 
     checkpoint_best_only: bool = False
@@ -195,7 +195,7 @@ class TrainConfig:
     plots that are logged or saved at specified intervals.
     """
 
-    _subfolders: list = field(default_factory=list)
+    _subfolders: list[str] = field(default_factory=list)
     """List of subfolders used for logging different runs using the same config inside the.
 
     root folder.
@@ -203,6 +203,67 @@ class TrainConfig:
     Each subfolder is of structure `<id>_<timestamp>_<PID>`.
     """
 
+    nprocs: int = 1
+    """
+    The number of processes to use for training when spawning subprocesses.
+
+    For effective parallel processing, set this to a value greater than 1.
+    - In case of Multi-GPU or Multi-Node-Multi-GPU setups, nprocs should be equal to
+    the total number of GPUs across all nodes (world size), or total number of GPU to be used.
+
+    If nprocs > 1, multiple processes will be spawned for training. The training framework will launch
+    additional processes (e.g., for distributed or parallel training).
+    - For CPU setup, this will launch a true parallel processes
+    - For GPU setup, this will launch a distributed training routine.
+    This uses the DistributedDataParallel framework from PyTorch.
+    """
+
+    compute_setup: str = "cpu"
+    """
+    Compute device setup; options are "auto", "gpu", or "cpu".
+
+    - "auto": Automatically uses GPU if available; otherwise, falls back to CPU.
+    - "gpu": Forces GPU usage, raising an error if no CUDA device is available.
+    - "cpu": Forces the use of CPU regardless of GPU availability.
+    """
+
+    backend: str = "gloo"
+    """
+    Backend used for distributed training communication.
+
+    The default is "gloo". Other options may include "nccl" - which is optimized for GPU-based training or "mpi",
+    depending on your system and requirements.
+    It should be one of the backends supported by `torch.distributed`. For further details, please look at
+    [torch backends](https://pytorch.org/docs/stable/distributed.html#torch.distributed.Backend)
+    """
+
+    log_setup: str = "cpu"
+    """
+    Logging device setup; options are "auto" or "cpu".
+
+    - "auto": Uses the same device for logging as for computation.
+    - "cpu": Forces logging to occur on the CPU. This can be useful to avoid potential conflicts with GPU processes.
+    """
+
+    dtype: dtype | None = None
+    """
+    Data type (precision) for computations.
+
+    Both model parameters, and dataset will be of the provided precision.
+
+    If not specified or None, the default torch precision (usually torch.float32) is used.
+    If provided dtype is torch.complex128, model parameters will be torch.complex128, and data parameters will be torch.float64
+    """
+
+    all_reduce_metrics: bool = False
+    """
+    Whether to aggregate metrics (e.g., loss, accuracy) across processes.
+
+    When True, metrics from different training processes are averaged to provide a consolidated metrics.
+    Note: Since aggregation requires synchronization/all_reduce operation, this can increase the
+     computation time significantly.
+    """
+
 
 @dataclass
 class FeatureMapConfig: