qadence 1.10.1__tar.gz → 1.10.3__tar.gz

{qadence-1.10.1 → qadence-1.10.3}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: qadence
-Version: 1.10.1
+Version: 1.10.3
 Summary: Pasqal interface for circuit-based quantum computing SDKs
-Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>, Giorgio Tosti Balducci <giorgio.tosti-balducci@pasqal.com>, Daniele Cucurachi <daniele.cucurachi@pasqal.com>, Pim Venderbosch <pim.venderbosch@pasqal.com>
+Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>, Giorgio Tosti Balducci <giorgio.tosti-balducci@pasqal.com>, Daniele Cucurachi <daniele.cucurachi@pasqal.com>, Pim Venderbosch <pim.venderbosch@pasqal.com>, Manu Lahariya <manu.lahariya@pasqal.com>
 License: Apache 2.0
 License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
@@ -43,7 +43,7 @@ Requires-Dist: nvidia-pyindex; extra == 'dlprof'
 Provides-Extra: horqrux
 Requires-Dist: einops; extra == 'horqrux'
 Requires-Dist: flax; extra == 'horqrux'
-Requires-Dist: horqrux==0.6.2; extra == 'horqrux'
+Requires-Dist: horqrux==0.7.0; extra == 'horqrux'
 Requires-Dist: jax; extra == 'horqrux'
 Requires-Dist: jaxopt; extra == 'horqrux'
 Requires-Dist: optax; extra == 'horqrux'
@@ -56,8 +56,8 @@ Provides-Extra: protocols
 Requires-Dist: qadence-protocols; extra == 'protocols'
 Provides-Extra: pulser
 Requires-Dist: pasqal-cloud==0.12.7; extra == 'pulser'
-Requires-Dist: pulser-core==1.2.0; extra == 'pulser'
-Requires-Dist: pulser-simulation==1.2.0; extra == 'pulser'
+Requires-Dist: pulser-core==1.2.2; extra == 'pulser'
+Requires-Dist: pulser-simulation==1.2.2; extra == 'pulser'
 Provides-Extra: visualization
 Requires-Dist: graphviz; extra == 'visualization'
 Description-Content-Type: text/markdown
@@ -75,7 +75,10 @@ programs** with tunable qubit interactions and arbitrary register topologies rea
 
 **For a high-level overview of Qadence features, [check out our white paper](https://arxiv.org/abs/2401.09915).**
 
-**For more detailed information, [check out the documentation](https://pasqal-io.github.io/qadence/latest/).**
+**For more detailed information, [check out the documentation](https://pasqal-io.github.io/qadence/latest/).
+
+**For any questions or comments, [feel free to start a discussion](https://github.com/pasqal-io/qadence/discussions).
+**
 
 [![Linting](https://github.com/pasqal-io/qadence/actions/workflows/lint.yml/badge.svg)](https://github.com/pasqal-io/qadence/actions/workflows/lint.yml)
 [![Tests](https://github.com/pasqal-io/qadence/actions/workflows/test_fast.yml/badge.svg)](https://github.com/pasqal-io/qadence/actions/workflows/test_fast.yml)

{qadence-1.10.1 → qadence-1.10.3}/README.md

@@ -11,7 +11,10 @@ programs** with tunable qubit interactions and arbitrary register topologies rea
 
 **For a high-level overview of Qadence features, [check out our white paper](https://arxiv.org/abs/2401.09915).**
 
-**For more detailed information, [check out the documentation](https://pasqal-io.github.io/qadence/latest/).**
+**For more detailed information, [check out the documentation](https://pasqal-io.github.io/qadence/latest/).
+
+**For any questions or comments, [feel free to start a discussion](https://github.com/pasqal-io/qadence/discussions).
+**
 
 [![Linting](https://github.com/pasqal-io/qadence/actions/workflows/lint.yml/badge.svg)](https://github.com/pasqal-io/qadence/actions/workflows/lint.yml)
 [![Tests](https://github.com/pasqal-io/qadence/actions/workflows/test_fast.yml/badge.svg)](https://github.com/pasqal-io/qadence/actions/workflows/test_fast.yml)

{qadence-1.10.1 → qadence-1.10.3}/mkdocs.yml

@@ -19,6 +19,7 @@ nav:
     - Quantum models: content/quantummodels.md
     - Quantum registers: content/register.md
     - State initialization: content/state_init.md
+    - Noisy Simulation: content/noisy_simulation.md
     - Arbitrary Hamiltonians: content/hamiltonians.md
     - Time-dependent generators: content/time_dependent.md
     - QML Constructors: content/qml_constructors.md
@@ -92,6 +93,8 @@ nav:
       - Pulser: api/backends/pulser.md
       - DifferentiableBackend: api/backends/differentiable.md
 
+  - Contact: https://github.com/pasqal-io/qadence/discussions
+
 edit_uri: edit/main/docs/
 
 theme:

{qadence-1.10.1 → qadence-1.10.3}/pyproject.toml

@@ -23,11 +23,12 @@ authors = [
     { name = "Charles Moussa", email = "charles.moussa@pasqal.com" },
     { name = "Giorgio Tosti Balducci", email = "giorgio.tosti-balducci@pasqal.com" },
     { name = "Daniele Cucurachi", email = "daniele.cucurachi@pasqal.com" },
-    { name = "Pim Venderbosch", email = "pim.venderbosch@pasqal.com" }
+    { name = "Pim Venderbosch", email = "pim.venderbosch@pasqal.com" },
+    { name = "Manu Lahariya", email = "manu.lahariya@pasqal.com" },
 ]
 requires-python = ">=3.9"
 license = { text = "Apache 2.0" }
-version = "1.10.1"
+version = "1.10.3"
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python",
@@ -64,8 +65,8 @@ allow-ambiguous-features = true
 
 [project.optional-dependencies]
 pulser = [
-    "pulser-core==1.2.0",
-    "pulser-simulation==1.2.0",
+    "pulser-core==1.2.2",
+    "pulser-simulation==1.2.2",
     "pasqal-cloud==0.12.7",
 ]
 visualization = [
@@ -75,7 +76,7 @@ visualization = [
     # "scour",
 ]
 horqrux = [
-    "horqrux==0.6.2",
+    "horqrux==0.7.0",
     "jax",
     "flax",
     "optax",
@@ -136,7 +137,7 @@ filterwarnings = [
 [tool.hatch.envs.docs]
 dependencies = [
     "mkdocs",
-    "mkdocs_autorefs<1.3.1",
+    "mkdocs_autorefs",
     "mkdocs-material",
     "mkdocstrings",
     "mkdocstrings-python",

{qadence-1.10.1 → qadence-1.10.3}/qadence/backends/horqrux/convert_ops.py

@@ -245,7 +245,7 @@ class HorqHamiltonianEvolution(NativeHorqHEvo):
 
         self._time_evolution = lambda values: values[self.param_names[0]]
 
-    def unitary(self, values: dict[str, Array]) -> Array:
+    def _unitary(self, values: dict[str, Array]) -> Array:
         """The evolved operator given current parameter values for generator and time evolution."""
         return expm(self._hamiltonian(self, values) * (-1j * self._time_evolution(values)))
 

{qadence-1.10.1 → qadence-1.10.3}/qadence/blocks/block_to_tensor.py

@@ -79,6 +79,7 @@ def _fill_identities(
         torch.Tensor: augmented matrix with dimensions (2**nqubits, 2**nqubits)
         or a tensor (2**n_qubits) if diag_only
     """
+    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
@@ -469,14 +470,13 @@ def _block_to_tensor_embedded(
         )
 
     elif isinstance(block, MatrixBlock):
-        mat = block.matrix.unsqueeze(0)
-        # FIXME: properly handle identity filling in matrix blocks
-        # mat = _fill_identities(
-        #    block.matrix.unsqueeze(0),
-        #    block.qubit_support,
-        #    qubit_support,
-        #    endianness=endianness,
-        # )
+        mat = _fill_identities(
+            block.matrix.unsqueeze(0),
+            block.qubit_support,
+            qubit_support,
+            endianness=endianness,
+            device=device,
+        )
 
     elif isinstance(block, SWAP):
         swap_block = _swap_block(block)

{qadence-1.10.1 → qadence-1.10.3}/qadence/blocks/matrix.py

@@ -7,6 +7,8 @@ import numpy as np
 import torch
 from torch.linalg import eigvals
 
+from math import log
+
 from qadence.blocks import PrimitiveBlock
 from qadence.noise import NoiseHandler
 
@@ -84,6 +86,8 @@ class MatrixBlock(PrimitiveBlock):
             if not self.is_unitary(matrix):
                 logger.warning("Provided matrix is not unitary.")
         self.matrix = matrix.clone()
+        if int(log(self.matrix.size(1), 2)) != len(qubit_support):
+            raise ValueError("Provided matrix does not match the qubit_support length.")
         super().__init__(qubit_support, noise)
 
     @cached_property

{qadence-1.10.1 → qadence-1.10.3}/qadence/ml_tools/__init__.py

@@ -4,6 +4,7 @@ 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 .data import DictDataLoader, InfiniteTensorDataset, OptimizeResult, to_dataloader
+from .information import InformationContent
 from .models import QNN
 from .optimize_step import optimize_step as default_optimize_step
 from .parameters import get_parameters, num_parameters, set_parameters

{qadence-1.10.1 → qadence-1.10.3}/qadence/ml_tools/config.py

@@ -434,15 +434,17 @@ class AnsatzConfig:
     """What type of ansatz.
 
     `AnsatzType.HEA` for Hardware Efficient Ansatz.
-    `AnsatzType.IIA` for Identity intialized Ansatz.
+    `AnsatzType.IIA` for Identity Intialized Ansatz.
+    `AnsatzType.ALA` for Alternating Layer Ansatz.
     """
 
     ansatz_strategy: Strategy = Strategy.DIGITAL
     """Ansatz strategy.
 
-    `Strategy.DIGITAL` for fully digital ansatz. Required if `ansatz_type` is `AnsatzType.IIA`.
-    `Strategy.SDAQC` for analog entangling block.
-    `Strategy.RYDBERG` for fully rydberg hea ansatz.
+    `Strategy.DIGITAL` for fully digital ansatz. Required if `ansatz_type` is `AnsatzType.ALA`.
+    `Strategy.SDAQC` for analog entangling block. Only available for `AnsatzType.HEA` or
+    `AnsatzType.ALA`.
+    `Strategy.RYDBERG` for fully rydberg hea ansatz. Only available for `AnsatzType.HEA`.
     """
 
     strategy_args: dict = field(default_factory=dict)
@@ -484,6 +486,13 @@ class AnsatzConfig:
     """
     # The default for a dataclass can not be a mutable object without using this default_factory.
 
+    m_block_qubits: int | None = None
+    """
+    The number of qubits in the local entangling block of an Alternating Layer Ansatz (ALA).
+
+    Only used when `ansatz_type` is `AnsatzType.ALA`.
+    """
+
     param_prefix: str = "theta"
     """The base bame of the variational parameter."""
 
@@ -499,3 +508,13 @@ class AnsatzConfig:
             assert (
                 self.ansatz_strategy != Strategy.RYDBERG
             ), "Rydberg strategy not allowed for Identity-initialized ansatz."
+
+        if self.ansatz_type == AnsatzType.ALA:
+            assert (
+                self.ansatz_strategy == Strategy.DIGITAL
+            ), f"{self.ansatz_strategy} not allowed for Alternating Layer Ansatz.\
+            Only `Strategy.DIGITAL` allowed."
+
+            assert (
+                self.m_block_qubits is not None
+            ), "m_block_qubits must be specified for Alternating Layer Ansatz."

{qadence-1.10.1 → qadence-1.10.3}/qadence/ml_tools/constructors.py

@@ -13,13 +13,14 @@ from qadence.constructors import (
     analog_feature_map,
     feature_map,
     hamiltonian_factory,
-    iia,
     rydberg_feature_map,
     rydberg_hea,
     rydberg_tower_feature_map,
 )
+from qadence.constructors.ala import ala_digital
 from qadence.constructors.hamiltonians import ObservableConfig, TDetuning
 from qadence.constructors.hea import hea_digital, hea_sDAQC
+from qadence.constructors.iia import iia
 from qadence.measurements import Measurements
 from qadence.noise import NoiseHandler
 from qadence.operations import CNOT, RX, RY, I, N, Z
@@ -596,6 +597,58 @@ def _create_hea(
         )
 
 
+def _create_ala_digital(
+    num_qubits: int,
+    config: AnsatzConfig,
+) -> AbstractBlock:
+    """
+    Create the Digital Alternating Layer Ansatz based on the configuration.
+
+    Args:
+        num_qubits (int): The number of qubits.
+        config (AnsatzConfig): The configuration for the ansatz.
+
+    Returns:
+        AbstractBlock: The Digital Alternating Layer Ansatz.
+    """
+    operations = config.strategy_args.get("operation", [RX, RY, RX])
+    entangler = config.strategy_args.get("entangler", CNOT)
+
+    return ala_digital(
+        n_qubits=num_qubits,
+        m_block_qubits=config.m_block_qubits,  # type: ignore[arg-type]
+        param_prefix=config.param_prefix,
+        operations=operations,
+        entangler=entangler,
+    )
+
+
+def _create_ala(
+    num_qubits: int,
+    config: AnsatzConfig,
+) -> AbstractBlock:
+    """
+    Create the Alternating Layer Ansatz based on the configuration.
+
+    Args:
+        num_qubits (int): The number of qubits.
+        config (AnsatzConfig): The configuration for the ansatz.
+
+    Returns:
+        AbstractBlock: The Alternating Layer Ansatz.
+
+    Raises:
+        ValueError: If the ansatz strategy is not `Strategy.DIGITAL`.
+    """
+    if config.ansatz_strategy == Strategy.DIGITAL:
+        return _create_ala_digital(num_qubits=num_qubits, config=config)
+    else:
+        raise ValueError(
+            f"Invalid ansatz strategy {config.ansatz_strategy} provided. Only `Strategy.DIGITAL` \
+                allowed"
+        )
+
+
 def create_ansatz(
     register: int | Register,
     config: AnsatzConfig,
@@ -619,6 +672,8 @@ def create_ansatz(
         return _create_iia(num_qubits=num_qubits, config=config)
     elif config.ansatz_type == AnsatzType.HEA:
         return _create_hea(register=register, config=config)
+    elif config.ansatz_type == AnsatzType.ALA:
+        return _create_ala(num_qubits=num_qubits, config=config)
     else:
         raise NotImplementedError(
             f"Ansatz of type {config.ansatz_type} not implemented yet. Only `AnsatzType.HEA` and\

qadence-1.10.3/qadence/ml_tools/information/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+from .information_content import InformationContent

qadence-1.10.3/qadence/ml_tools/information/information_content.py

@@ -0,0 +1,339 @@
+from __future__ import annotations
+
+import functools
+from logging import getLogger
+from math import log, sqrt
+from statistics import NormalDist
+from typing import Any, Callable
+
+import torch
+from torch import nn
+from torch.func import functional_call  # type: ignore
+
+logger = getLogger("ml_tools")
+
+
+class InformationContent:
+    def __init__(
+        self,
+        model: nn.Module,
+        loss_fn: Callable,
+        xs: Any,
+        epsilons: torch.Tensor,
+        variation_multiple: int = 20,
+    ) -> None:
+        """Information Landscape class.
+
+        This class handles the study of loss landscape from information theoretic
+        perspective and provides methods to get bounds on the norm of the
+        gradient from the Information Content of the loss landscape.
+
+        Args:
+            model: The quantum or classical model to analyze.
+            loss_fn: Loss function that takes model output and calculates loss
+            xs: Input data to evaluate the model on
+            epsilons: The thresholds to use for discretization of the finite derivatives
+            variation_multiple: The number of sets of variational parameters to generate per each
+                variational parameter. The number of variational parameters required for the
+                statistical analysis scales linearly with the amount of them present in the
+                model. This is that linear factor.
+
+        Notes:
+            This class provides flexibility in terms of what the model, the loss function,
+            and the xs are. The only requirement is that the loss_fn takes the model and xs as
+            arguments and returns the loss, and another dictionary of other metrics.
+
+            Thus, assumed structure:
+                loss_fn(model, xs) -> (loss, metrics, ...)
+
+            Example: A Classifier
+                ```python
+                model = nn.Linear(10, 1)
+
+                def loss_fn(
+                    model: nn.Module,
+                    xs: tuple[torch.Tensor, torch.Tensor]
+                ) -> tuple[torch.Tensor, dict[str, float]:
+                    criterion = nn.MSELoss()
+                    inputs, labels = xs
+                    outputs = model(inputs)
+                    loss = criterion(outputs, labels)
+                    metrics = {"loss": loss.item()}
+                    return loss, metrics
+
+                xs = (torch.randn(10, 10), torch.randn(10, 1))
+
+                info_landscape = InfoLandscape(model, loss_fn, xs)
+                ```
+                In this example, the model is a linear classifier, and the `xs` include both the
+                inputs and the target labels. The logic for calculation of the loss from this lies
+                entirely within the `loss_fn` function. This can then further be used to obtain the
+                bounds on the average norm of the gradient of the loss function.
+
+            Example: A Physics Informed Neural Network
+                ```python
+                class PhysicsInformedNN(nn.Module):
+                    // <Initialization Logic>
+
+                    def forward(self, xs: dict[str, torch.Tensor]):
+                        return {
+                            "pde_residual": pde_residual(xs["pde"]),
+                            "boundary_condition": bc_term(xs["bc"]),
+                        }
+
+                def loss_fn(
+                    model: PhysicsInformedNN,
+                    xs: dict[str, torch.Tensor]
+                ) -> tuple[torch.Tensor, dict[str, float]:
+                    pde_residual, bc_term = model(xs)
+                    loss = torch.mean(torch.sum(pde_residual**2, dim=1), dim=0)
+                        + torch.mean(torch.sum(bc_term**2, dim=1), dim=0)
+
+                    return loss, {"pde_residual": pde_residual, "bc_term": bc_term}
+
+                xs = {
+                    "pde": torch.linspace(0, 1, 10),
+                    "bc": torch.tensor([0.0]),
+                }
+
+                info_landscape = InfoLandscape(model, loss_fn, xs)
+                ```
+
+                In this example, the model is a Physics Informed Neural Network, and the `xs`
+                are the inputs to the different residual components of the model. The logic
+                for calculation of the residuals lies within the PhysicsInformedNN class, and
+                the loss function is defined to calculate the loss that is to be optimized
+                from these residuals. This can then further be used to obtain the
+                bounds on the average norm of the gradient of the loss function.
+
+            The first value that the `loss_fn` returns is the loss value that is being optimized.
+            The function is also expected to return other value(s), often the metrics that are
+            used to calculate the loss. These values are ignored for the purpose of this class.
+        """
+        self.model = model
+        self.loss_fn = loss_fn
+        self.xs = xs
+        self.epsilons = epsilons
+        self.device = next(model.parameters()).device
+
+        self.param_shapes = {}
+        self.total_params = 0
+
+        for name, param in model.named_parameters():
+            self.param_shapes[name] = param.shape
+            self.total_params += param.numel()
+        self.n_variations = variation_multiple * self.total_params
+        self.all_variations = torch.empty(
+            (self.n_variations, self.total_params), device=self.device
+        ).uniform_(0, 2 * torch.pi)
+
+    def reshape_param_variations(self) -> dict[str, torch.Tensor]:
+        """Reshape variations of the model's variational parameters.
+
+        Returns:
+            Dictionary of parameter tensors, each with shape [n_variations, *param_shape]
+        """
+        param_variations = {}
+        start_idx = 0
+
+        for name, shape in self.param_shapes.items():
+            param_size = torch.prod(torch.tensor(shape)).item()
+            param_variations[name] = self.all_variations[
+                :, start_idx : start_idx + param_size
+            ].view(self.n_variations, *shape)
+            start_idx += param_size
+
+        return param_variations
+
+    def batched_loss(self) -> torch.Tensor:
+        """Calculate loss for all parameter variations in a batched manner.
+
+        Returns: Tensor of loss values for each parameter variation
+        """
+        param_variations = self.reshape_param_variations()
+        losses = torch.zeros(self.n_variations, device=self.device)
+
+        for i in range(self.n_variations):
+            params = {name: param[i] for name, param in param_variations.items()}
+            current_model = lambda x: functional_call(self.model, params, (x,))
+            losses[i] = self.loss_fn(current_model, self.xs)[0]
+
+        return losses
+
+    def randomized_finite_der(self) -> torch.Tensor:
+        """
+        Calculate normalized finite difference of loss on doing random walk in the parameter space.
+
+        This serves as a proxy for the derivative of the loss with respect to parameters.
+
+        Returns:
+            Tensor containing normalized finite differences (approximate directional derivatives)
+            between consecutive points in the random walk. Shape: [n_variations - 1]
+        """
+        losses = self.batched_loss()
+
+        return (losses[1:] - losses[:-1]) / (
+            torch.norm(self.all_variations[1:] - self.all_variations[:-1], dim=1) + 1e-8
+        )
+
+    def discretize_derivatives(self) -> torch.Tensor:
+        """
+        Convert finite derivatives into discrete values.
+
+        Returns:
+            Tensor containing discretized derivatives with shape [n_epsilons, n_variations-2]
+            Each row contains {-1, 0, 1} values for that epsilon
+        """
+        derivatives = self.randomized_finite_der()
+
+        derivatives = derivatives.unsqueeze(0)
+        epsilons = self.epsilons.unsqueeze(1)
+
+        discretized = torch.zeros((len(epsilons), len(derivatives[0])), device=self.device)
+        discretized[derivatives > epsilons] = 1
+        discretized[derivatives < -epsilons] = -1
+
+        return discretized
+
+    def calculate_transition_probabilities_batch(self) -> torch.Tensor:
+        """
+        Calculate transition probabilities for multiple epsilon values.
+
+        Returns:
+            Tensor of shape [n_epsilons, 6] containing probabilities for each transition type
+            Columns order: [+1to0, +1to-1, 0to+1, 0to-1, -1to0, -1to+1]
+        """
+        discretized = self.discretize_derivatives()
+
+        current = discretized[:, :-1]
+        next_val = discretized[:, 1:]
+
+        transitions = torch.stack(
+            [
+                ((current == 1) & (next_val == 0)).sum(dim=1),
+                ((current == 1) & (next_val == -1)).sum(dim=1),
+                ((current == 0) & (next_val == 1)).sum(dim=1),
+                ((current == 0) & (next_val == -1)).sum(dim=1),
+                ((current == -1) & (next_val == 0)).sum(dim=1),
+                ((current == -1) & (next_val == 1)).sum(dim=1),
+            ],
+            dim=1,
+        ).float()
+
+        total_transitions = current.size(1)
+        probabilities = transitions / total_transitions
+
+        return probabilities
+
+    @functools.cached_property
+    def calculate_IC(self) -> torch.Tensor:
+        """
+        Calculate Information Content for multiple epsilon values.
+
+        Returns: Tensor of IC values for each epsilon [n_epsilons]
+        """
+        probs = self.calculate_transition_probabilities_batch()
+
+        mask = probs > 1e-4
+
+        ic_terms = torch.where(mask, -probs * torch.log(probs), torch.zeros_like(probs))
+        ic_values = ic_terms.sum(dim=1) / torch.log(torch.tensor(6.0))
+
+        return ic_values
+
+    def max_IC(self) -> tuple[float, float]:
+        """
+        Get the maximum Information Content and its corresponding epsilon.
+
+        Returns: Tuple of (maximum IC value, optimal epsilon)
+        """
+        max_ic, max_idx = torch.max(self.calculate_IC, dim=0)
+        max_epsilon = self.epsilons[max_idx]
+        return max_ic.item(), max_epsilon.item()
+
+    def sensitivity_IC(self, eta: float) -> float:
+        """
+        Find the minimum value of epsilon such that the information content is less than eta.
+
+        Args:
+            eta: Threshold value, the sensitivity IC.
+
+        Returns: The epsilon value that gives IC that is less than the sensitivity IC.
+        """
+        ic_values = self.calculate_IC
+        mask = ic_values < eta
+        epsilons = self.epsilons[mask]
+        return float(epsilons.min().item())
+
+    @staticmethod
+    @functools.lru_cache
+    def q_value(H_value: float) -> float:
+        """
+        Compute the q value.
+
+        q is the solution to the equation:
+        H(x) = 4h(x) + 2h(1/2 - 2x)
+
+        It is the value of the probability of 4 of the 6 transitions such that
+        the IC is the same as the IC of our system.
+
+        This quantity is useful in calculating the bounds on the norms of the gradients.
+
+        Args:
+            H_value (float): The information content.
+
+        Returns:
+            float: The q value
+        """
+
+        x = torch.linspace(0.001, 0.16667, 10000)
+
+        H = -4 * x * torch.log(x) / torch.log(torch.tensor(6)) - 2 * (0.5 - 2 * x) * torch.log(
+            0.5 - 2 * x
+        ) / torch.log(torch.tensor(6))
+        err = torch.abs(H - H_value)
+        idx = torch.argmin(err)
+        return float(x[idx].item())
+
+    def get_grad_norm_bounds_max_IC(self) -> tuple[float, float]:
+        """
+        Compute the bounds on the average norm of the gradient.
+
+        Returns:
+            tuple[Tensor, Tensor]: The lower and upper bounds.
+        """
+        max_IC, epsilon_m = self.max_IC()
+        lower_bound = (
+            epsilon_m
+            * sqrt(self.total_params)
+            / (NormalDist().inv_cdf(1 - 2 * self.q_value(max_IC)))
+        )
+        upper_bound = (
+            epsilon_m
+            * sqrt(self.total_params)
+            / (NormalDist().inv_cdf(0.5 * (1 + 2 * self.q_value(max_IC))))
+        )
+
+        if max_IC < log(2, 6):
+            logger.warning(
+                "Warning: The maximum IC is less than the required value. The bounds may be"
+                + " inaccurate."
+            )
+
+        return lower_bound, upper_bound
+
+    def get_grad_norm_bounds_sensitivity_IC(self, eta: float) -> float:
+        """
+        Compute the bounds on the average norm of the gradient.
+
+        Args:
+            eta (float): The sensitivity IC.
+
+        Returns:
+            Tensor: The lower bound.
+        """
+        epsilon_sensitivity = self.sensitivity_IC(eta)
+        upper_bound = (
+            epsilon_sensitivity * sqrt(self.total_params) / (NormalDist().inv_cdf(1 - 3 * eta / 2))
+        )
+        return upper_bound