From 5062c1e1545c2e62a61b0022fdb2533c1a439b1b Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Mon, 11 Nov 2024 11:00:41 +0100
Subject: [PATCH 01/15] adding analog noise

---
 pyqtorch/hamiltonians/evolution.py | 46 +++++++++++++++++++++-------
 pyqtorch/noise/gates.py            | 16 ++++++----
 tests/test_analog.py               | 48 ++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 15 deletions(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index d819a128..a8985b03 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -16,6 +16,7 @@
 from pyqtorch.embed import ConcretizedCallable, Embedding
 from pyqtorch.primitives import Primitive
 from pyqtorch.quantum_operation import QuantumOperation
+from pyqtorch.time_dependent.mesolve import mesolve
 from pyqtorch.time_dependent.sesolve import sesolve
 from pyqtorch.utils import (
     ATOL,
@@ -140,6 +141,11 @@ class HamiltonianEvolution(Sequence):
         operations: List of operations.
         cache_length: LRU cache cache_length evolution operators for given set
                     of parameter values.
+        duration: Total duration for evolving when using a solver.
+        steps: Number of steps to use when using solver.
+        solver: Time-dependent Lindblad master equation solver.
+        noise_operators: List of tensors or Kraus oeprators adding analog noise
+            when solving with a Shrodinger equation solver.
     """
 
     def __init__(
@@ -151,6 +157,7 @@ def __init__(
         duration: Tensor | str | float | None = None,
         steps: int = 100,
         solver=SolverType.DP5_SE,
+        noise_operators: list[Tensor] = list(),
     ):
         """Initializes the HamiltonianEvolution.
         Depending on the generator argument, set the type and set the right generator getter.
@@ -161,7 +168,13 @@ def __init__(
             qubit_support: The qubits the operator acts on. If generator is a quantum
                 operation or sequence of operations,
                 it will be inferred from the generator.
-            generator_parametric: Whether the generator is parametric or not.
+            cache_length: LRU cache cache_length evolution operators for given set
+                    of parameter values.
+            duration: Total duration for evolving when using a solver.
+            steps: Number of steps to use when using solver.
+            solver: Time-dependent Lindblad master equation solver.
+            noise_operators: List of tensors or Kraus oeprators adding analog noise
+                when solving with a Shrodinger equation solver.
         """
 
         self.solver_type = solver
@@ -243,6 +256,8 @@ def __init__(
         self._cache_hamiltonian_evo: dict[str, Tensor] = dict()
         self.cache_length = cache_length
 
+        self.noise_operators: list[Tensor] = noise_operators
+
     @property
     def generator(self) -> ModuleList:
         """Returns the operations making the generator.
@@ -411,17 +426,28 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 .squeeze(2)
             )
 
-        sol = sesolve(
-            Ht,
-            torch.flatten(state, start_dim=0, end_dim=-2),
-            t_grid,
-            self.solver_type,
-        )
+        if len(self.noise_operators) == 0:
+            sol = sesolve(
+                Ht,
+                torch.flatten(state, start_dim=0, end_dim=-2),
+                t_grid,
+                self.solver_type,
+            )
 
-        # Retrieve the last state of shape (2**n_qubits, batch_size)
-        state = sol.states[-1]
+            # Retrieve the last state of shape (2**n_qubits, batch_size)
+            state = sol.states[-1]
 
-        return state.reshape([2] * n_qubits + [batch_size])
+            return state.reshape([2] * n_qubits + [batch_size])
+        else:
+            sol = mesolve(
+                Ht,
+                torch.flatten(state, start_dim=0, end_dim=-2),
+                self.noise_operators,
+                t_grid,
+                self.solver_type,
+            )
+            state = sol.states[-1]
+            return state.reshape([2] * n_qubits * 2 + [batch_size])
 
     def forward(
         self,
diff --git a/pyqtorch/noise/gates.py b/pyqtorch/noise/gates.py
index 19a9c7b4..136c0aa9 100644
--- a/pyqtorch/noise/gates.py
+++ b/pyqtorch/noise/gates.py
@@ -9,7 +9,12 @@
 from pyqtorch.apply import apply_operator_dm
 from pyqtorch.embed import Embedding
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, XMAT, YMAT, ZMAT
-from pyqtorch.utils import DensityMatrix, density_mat, qubit_support_as_tuple
+from pyqtorch.utils import (
+    DensityMatrix,
+    density_mat,
+    promote_operator,
+    qubit_support_as_tuple,
+)
 
 
 class Noise(torch.nn.Module):
@@ -35,11 +40,12 @@ def extra_repr(self) -> str:
     def kraus_operators(self) -> list[Tensor]:
         return [getattr(self, f"kraus_{i}") for i in range(len(self._buffers))]
 
-    def tensor(
-        self,
-    ) -> list[Tensor]:
+    def tensor(self, n_qubit_support: int | None = None) -> list[Tensor]:
         # Since PyQ expects tensor.Size = [2**n_qubits, 2**n_qubits,batch_size].
-        return [kraus_op.unsqueeze(2) for kraus_op in self.kraus_operators]
+        t_ops = [kraus_op.unsqueeze(2) for kraus_op in self.kraus_operators]
+        if n_qubit_support is None:
+            return t_ops
+        return [promote_operator(t, self.target, n_qubit_support) for t in t_ops]
 
     def forward(
         self,
diff --git a/tests/test_analog.py b/tests/test_analog.py
index 64ac0d3d..436a4a5d 100644
--- a/tests/test_analog.py
+++ b/tests/test_analog.py
@@ -18,10 +18,12 @@
     XMAT,
     ZMAT,
 )
+from pyqtorch.noise import Depolarizing
 from pyqtorch.utils import (
     ATOL,
     RTOL,
     SolverType,
+    density_mat,
     is_normalized,
     operator_kron,
     overlap,
@@ -312,6 +314,7 @@ def test_timedependent(
 
     # simulate with time-dependent solver
     t_points = torch.linspace(0, dur_val[0], n_steps)
+
     psi_solver = pyq.sesolve(
         torch_hamiltonian, psi_start.reshape(-1, 1), t_points, ode_solver
     ).states[-1]
@@ -336,6 +339,51 @@ def test_timedependent(
     assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)
 
 
+@pytest.mark.parametrize("duration", [torch.rand(1), "duration"])
+def test_timedependent_with_noise(
+    tparam: str,
+    param_y: float,
+    duration: float,
+    n_steps: int,
+    torch_hamiltonian: Callable,
+    hamevo_generator: Sequence,
+    sin: tuple,
+    sq: tuple,
+) -> None:
+
+    psi_start = density_mat(random_state(2))
+    dur_val = duration if isinstance(duration, torch.Tensor) else torch.rand(1)
+
+    # simulate with time-dependent solver
+    t_points = torch.linspace(0, dur_val[0], n_steps)
+
+    list_ops = Depolarizing(0, error_probability=0.1).tensor(2)
+    solver = SolverType.DP5_ME
+    psi_solver = pyq.mesolve(
+        torch_hamiltonian, psi_start.reshape(-1, 1), list_ops, t_points, solver
+    ).states[-1]
+
+    # simulate with HamiltonianEvolution
+    embedding = pyq.Embedding(
+        tparam_name=tparam,
+        var_to_call={sin[0]: sin[1], sq[0]: sq[1]},
+    )
+    hamiltonian_evolution = pyq.HamiltonianEvolution(
+        generator=hamevo_generator,
+        time=tparam,
+        duration=duration,
+        steps=n_steps,
+        solver=solver,
+        noise_operators=list_ops,
+    )
+    values = {"y": param_y, "duration": dur_val}
+    psi_hamevo = hamiltonian_evolution(
+        state=psi_start, values=values, embedding=embedding
+    ).reshape(-1, 1)
+
+    assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)
+
+
 @pytest.mark.parametrize("n_qubits", [2, 4, 6])
 @pytest.mark.parametrize("batch_size", [1, 2])
 def test_hamevo_parametric_gen(n_qubits: int, batch_size: int) -> None:

From 9680d56b7b9ba13f822535c2c1760fe68203023f Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Mon, 11 Nov 2024 13:23:22 +0100
Subject: [PATCH 02/15] add test density_matrix

---
 pyqtorch/hamiltonians/evolution.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index a8985b03..0c94f8ab 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -29,6 +29,8 @@
     finitediff,
     is_diag,
     is_parametric,
+    DensityMatrix,
+    density_mat,
 )
 
 BATCH_DIM = 2
@@ -439,6 +441,8 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
 
             return state.reshape([2] * n_qubits + [batch_size])
         else:
+            if not isinstance(state, DensityMatrix):
+                state = density_mat(state)
             sol = mesolve(
                 Ht,
                 torch.flatten(state, start_dim=0, end_dim=-2),

From d1c41c8cf88f4b48aabbbdb7dda3b30fe279cb8d Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:03:29 +0100
Subject: [PATCH 03/15] add option use_sparse hamevo mesolve

---
 pyqtorch/hamiltonians/evolution.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index fa95ad83..4223878a 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -20,17 +20,17 @@
 from pyqtorch.time_dependent.sesolve import sesolve
 from pyqtorch.utils import (
     ATOL,
+    DensityMatrix,
     Operator,
     SolverType,
     State,
     StrEnum,
     _round_operator,
+    density_mat,
     expand_operator,
     finitediff,
     is_diag,
     is_parametric,
-    DensityMatrix,
-    density_mat,
 )
 
 BATCH_DIM = 2
@@ -437,7 +437,7 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 t_grid,
                 self.solver_type,
                 options={"use_sparse": self.use_sparse},
-        )
+            )
 
             # Retrieve the last state of shape (2**n_qubits, batch_size)
             state = sol.states[-1]
@@ -452,6 +452,7 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 self.noise_operators,
                 t_grid,
                 self.solver_type,
+                options={"use_sparse": self.use_sparse},
             )
             state = sol.states[-1]
             return state.reshape([2] * n_qubits * 2 + [batch_size])

From fe23d8fba03f54ab532100f32baf002143521f69 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:25:44 +0100
Subject: [PATCH 04/15] rm sesolve reshape

---
 pyqtorch/hamiltonians/evolution.py | 10 ++++------
 tests/test_analog.py               | 17 +++++++++++++----
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index 4223878a..a0003c6a 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -440,22 +440,20 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
             )
 
             # Retrieve the last state of shape (2**n_qubits, batch_size)
-            state = sol.states[-1]
-
-            return state.reshape([2] * n_qubits + [batch_size])
+            return sol.states[-1]
         else:
             if not isinstance(state, DensityMatrix):
                 state = density_mat(state)
             sol = mesolve(
                 Ht,
-                torch.flatten(state, start_dim=0, end_dim=-2),
+                state,
                 self.noise_operators,
                 t_grid,
                 self.solver_type,
                 options={"use_sparse": self.use_sparse},
             )
-            state = sol.states[-1]
-            return state.reshape([2] * n_qubits * 2 + [batch_size])
+            mesolve_state = sol.states[-1]
+            return mesolve_state
 
     def forward(
         self,
diff --git a/tests/test_analog.py b/tests/test_analog.py
index 4a6748d7..c6805f0b 100644
--- a/tests/test_analog.py
+++ b/tests/test_analog.py
@@ -349,10 +349,17 @@ def test_timedependent(
 
 
 @pytest.mark.parametrize("duration", [torch.rand(1), "duration"])
+@pytest.mark.parametrize(
+    "batchsize",
+    [
+        1,
+    ],
+)
 def test_timedependent_with_noise(
     tparam: str,
     param_y: float,
     duration: float,
+    batchsize: int,
     n_steps: int,
     torch_hamiltonian: Callable,
     hamevo_generator: Sequence,
@@ -360,16 +367,18 @@ def test_timedependent_with_noise(
     sq: tuple,
 ) -> None:
 
-    psi_start = density_mat(random_state(2))
+    psi_start = density_mat(random_state(2, batchsize))
     dur_val = duration if isinstance(duration, torch.Tensor) else torch.rand(1)
 
     # simulate with time-dependent solver
     t_points = torch.linspace(0, dur_val[0], n_steps)
 
+    # No batchsiwe for the jump operators
     list_ops = Depolarizing(0, error_probability=0.1).tensor(2)
+    list_ops = [op.squeeze() for op in list_ops]
     solver = SolverType.DP5_ME
     psi_solver = pyq.mesolve(
-        torch_hamiltonian, psi_start.reshape(-1, 1), list_ops, t_points, solver
+        torch_hamiltonian, psi_start, list_ops, t_points, solver
     ).states[-1]
 
     # simulate with HamiltonianEvolution
@@ -388,8 +397,8 @@ def test_timedependent_with_noise(
     values = {"y": param_y, "duration": dur_val}
     psi_hamevo = hamiltonian_evolution(
         state=psi_start, values=values, embedding=embedding
-    ).reshape(-1, 1)
-
+    )
+    print(psi_solver.shape, psi_hamevo.shape)
     assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)
 
 

From c5621873863dfefb7915e9d52c85398aa801d57d Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:28:28 +0100
Subject: [PATCH 05/15] fix reshape

---
 pyqtorch/hamiltonians/evolution.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index a0003c6a..45d8da80 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -440,7 +440,9 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
             )
 
             # Retrieve the last state of shape (2**n_qubits, batch_size)
-            return sol.states[-1]
+            state = sol.states[-1]
+
+            return state.reshape([2] * n_qubits + [batch_size])
         else:
             if not isinstance(state, DensityMatrix):
                 state = density_mat(state)

From 0ecbdfdcb8296e159413eb74c3a098edbaf7b027 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:31:38 +0100
Subject: [PATCH 06/15] return state for both noisy and non noisy

---
 pyqtorch/hamiltonians/evolution.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index 45d8da80..ff86f63f 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -440,9 +440,8 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
             )
 
             # Retrieve the last state of shape (2**n_qubits, batch_size)
-            state = sol.states[-1]
-
-            return state.reshape([2] * n_qubits + [batch_size])
+            # and reshape
+            state = sol.states[-1].reshape([2] * n_qubits + [batch_size])
         else:
             if not isinstance(state, DensityMatrix):
                 state = density_mat(state)
@@ -454,8 +453,9 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 self.solver_type,
                 options={"use_sparse": self.use_sparse},
             )
-            mesolve_state = sol.states[-1]
-            return mesolve_state
+            # Retrieve the last density matrix
+            state = sol.states[-1]
+        return state
 
     def forward(
         self,

From 64fa4eda3a7360f6e7516718e7e59c42a3e4d9d8 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:40:31 +0100
Subject: [PATCH 07/15] adding analog noise section

---
 docs/noise.md | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/docs/noise.md b/docs/noise.md
index 67e02972..bea18327 100644
--- a/docs/noise.md
+++ b/docs/noise.md
@@ -123,6 +123,46 @@ def fig_to_html(fig: Figure) -> str:  # markdown-exec: hide
 print(fig_to_html(plt.gcf())) # markdown-exec: hide
 ```
 
+## Analog Noise
+
+Analog noise is made possible by specifying `noise_operators` in `HamiltonianEvolution`:
+
+```python exec="on" source="material-block"
+import torch
+from pyqtorch import uniform_state, HamiltonianEvolution
+from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE
+from pyqtorch.noise import Depolarizing
+from pyqtorch.utils import  SolverType
+
+n_qubits = 2
+qubit_targets = list(range(n_qubits))
+
+# Random hermitian hamiltonian
+matrix = torch.rand(2**n_qubits, 2**n_qubits, dtype=DEFAULT_MATRIX_DTYPE)
+hermitian_matrix = matrix + matrix.T.conj()
+
+time = torch.tensor([1.0])
+time_symbol = "t"
+dur_val = torch.rand(1)
+list_ops = Depolarizing(0, error_probability=0.1).tensor(2)
+list_ops = [op.squeeze() for op in list_ops]
+solver = SolverType.DP5_ME
+n_steps = 5
+
+hamiltonian_evolution = HamiltonianEvolution(hermitian_matrix, time_symbol, qubit_targets,
+        duration=dur_val, steps=n_steps,
+        solver=solver, noise_operators=list_ops,)
+
+# Starting from a uniform state
+psi_start = uniform_state(n_qubits)
+
+# Returns the evolved state
+psi_end = hamiltonian_evolution(state = psi_start, values={time_symbol: time})
+
+print(psi_end)
+
+
+```
 
 ## Readout errors
 

From 7f4adb32b5e84314415d2d31c5f77d2910e60026 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:42:39 +0100
Subject: [PATCH 08/15] test batchsize and reshape necessary

---
 pyqtorch/hamiltonians/evolution.py | 2 +-
 tests/test_analog.py               | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index ff86f63f..35bd9dc4 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -454,7 +454,7 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 options={"use_sparse": self.use_sparse},
             )
             # Retrieve the last density matrix
-            state = sol.states[-1]
+            state = sol.states[-1].reshape(state.shape)
         return state
 
     def forward(
diff --git a/tests/test_analog.py b/tests/test_analog.py
index c6805f0b..8156e41a 100644
--- a/tests/test_analog.py
+++ b/tests/test_analog.py
@@ -353,6 +353,7 @@ def test_timedependent(
     "batchsize",
     [
         1,
+        3,
     ],
 )
 def test_timedependent_with_noise(

From 44fc7b5753b41bed67d6e29e1d7e13964baeb404 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:43:05 +0100
Subject: [PATCH 09/15] small comment

---
 pyqtorch/hamiltonians/evolution.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index 35bd9dc4..a243b3db 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -454,6 +454,7 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 options={"use_sparse": self.use_sparse},
             )
             # Retrieve the last density matrix
+            # and reshape
             state = sol.states[-1].reshape(state.shape)
         return state
 

From 0dc066c183e928a0edc2bfaff75fd3242277c254 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:44:34 +0100
Subject: [PATCH 10/15] fix docstring

---
 pyqtorch/hamiltonians/evolution.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index a243b3db..4ef6cf15 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -146,7 +146,7 @@ class HamiltonianEvolution(Sequence):
         duration: Total duration for evolving when using a solver.
         steps: Number of steps to use when using solver.
         solver: Time-dependent Lindblad master equation solver.
-        noise_operators: List of tensors or Kraus oeprators adding analog noise
+        noise_operators: List of tensors or Kraus operators adding analog noise
             when solving with a Shrodinger equation solver.
     """
 
@@ -176,7 +176,7 @@ def __init__(
             duration: Total duration for evolving when using a solver.
             steps: Number of steps to use when using solver.
             solver: Time-dependent Lindblad master equation solver.
-            noise_operators: List of tensors or Kraus oeprators adding analog noise
+            noise_operators: List of tensors or Kraus operators adding analog noise
                 when solving with a Shrodinger equation solver.
         """
 

From 816121f0bbdcf4107d021a89379239e6e3cc61d0 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 16:56:34 +0100
Subject: [PATCH 11/15] rm reshape

---
 pyqtorch/hamiltonians/evolution.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index 4ef6cf15..e41aa040 100644
--- a/pyqtorch/hamiltonians/evolution.py
+++ b/pyqtorch/hamiltonians/evolution.py
@@ -455,7 +455,7 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
             )
             # Retrieve the last density matrix
             # and reshape
-            state = sol.states[-1].reshape(state.shape)
+            state = sol.states[-1]
         return state
 
     def forward(

From 2e4349b471348fa90e45db427fc561505db7edb0 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 17:00:57 +0100
Subject: [PATCH 12/15] lower precision for mesolve

---
 tests/test_analog.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/test_analog.py b/tests/test_analog.py
index 8156e41a..ee35c106 100644
--- a/tests/test_analog.py
+++ b/tests/test_analog.py
@@ -399,8 +399,7 @@ def test_timedependent_with_noise(
     psi_hamevo = hamiltonian_evolution(
         state=psi_start, values=values, embedding=embedding
     )
-    print(psi_solver.shape, psi_hamevo.shape)
-    assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)
+    assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=1.0e-3)
 
 
 @pytest.mark.parametrize("n_qubits", [2, 4, 6])

From 245e6235e968daa4e3e20a59024c8a9b35ae89dd Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Fri, 15 Nov 2024 17:13:57 +0100
Subject: [PATCH 13/15] reducing atol for timedependent too

---
 tests/test_analog.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_analog.py b/tests/test_analog.py
index ee35c106..a253f6e9 100644
--- a/tests/test_analog.py
+++ b/tests/test_analog.py
@@ -345,7 +345,7 @@ def test_timedependent(
         state=psi_start, values=values, embedding=embedding
     ).reshape(-1, batch_size)
 
-    assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)
+    assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=1.0e-3)
 
 
 @pytest.mark.parametrize("duration", [torch.rand(1), "duration"])

From 6ca2d16384983b6f297ceed9cdbc4e1fdf3dfa32 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Mon, 18 Nov 2024 09:03:19 +0100
Subject: [PATCH 14/15] bump version

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

diff --git a/pyproject.toml b/pyproject.toml
index 23daf164..eab952b8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 name = "pyqtorch"
 description = "An efficient, large-scale emulator designed for quantum machine learning, seamlessly integrated with a PyTorch backend. Please refer to https://pyqtorch.readthedocs.io/en/latest/ for setup and usage info, along with the full documentation."
 readme = "README.md"
-version = "1.5.2"
+version = "1.5.3"
 requires-python = ">=3.8,<3.13"
 license = { text = "Apache 2.0" }
 keywords = ["quantum"]

From 86f87b32d9897441879350146cacedb32bd71d61 Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Wed, 20 Nov 2024 07:58:15 +0100
Subject: [PATCH 15/15] bump minor

---
 pyproject.toml       | 2 +-
 tests/test_analog.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index eab952b8..1ac6b19d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 name = "pyqtorch"
 description = "An efficient, large-scale emulator designed for quantum machine learning, seamlessly integrated with a PyTorch backend. Please refer to https://pyqtorch.readthedocs.io/en/latest/ for setup and usage info, along with the full documentation."
 readme = "README.md"
-version = "1.5.3"
+version = "1.6.0"
 requires-python = ">=3.8,<3.13"
 license = { text = "Apache 2.0" }
 keywords = ["quantum"]
diff --git a/tests/test_analog.py b/tests/test_analog.py
index a253f6e9..b21a8c18 100644
--- a/tests/test_analog.py
+++ b/tests/test_analog.py
@@ -374,7 +374,8 @@ def test_timedependent_with_noise(
     # simulate with time-dependent solver
     t_points = torch.linspace(0, dur_val[0], n_steps)
 
-    # No batchsiwe for the jump operators
+    # Define jump operators
+    # Note that we squeeze to remove the batch dimension
     list_ops = Depolarizing(0, error_probability=0.1).tensor(2)
     list_ops = [op.squeeze() for op in list_ops]
     solver = SolverType.DP5_ME