From 76abf6322c5e75b97c9a1d10ecc5ddad8b3f0f41 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Fri, 19 Jul 2024 19:27:38 +0200
Subject: [PATCH 01/55] JAX compatible givens decompositions. Unit tests
 passing.

---
 pennylane/qchem/givens_decomposition.py | 52 +++++++++++++------------
 1 file changed, 28 insertions(+), 24 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 18645237c98..ca7f6594e4f 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -16,7 +16,7 @@
 """
 
 import numpy as np
-
+import jax.numpy as jnp
 import pennylane as qml
 
 
@@ -41,21 +41,24 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         np.ndarray (or tensor): Givens rotation matrix
 
     """
-    abs_a, abs_b = np.abs(a), np.abs(b)
-    if abs_a < tol:
-        cosine, sine, phase = 1.0, 0.0, 1.0
-    elif abs_b < tol:
-        cosine, sine, phase = 0.0, 1.0, 1.0
-    else:
-        hypot = np.hypot(abs_a, abs_b)
-        cosine = abs_b / hypot
-        sine = abs_a / hypot
-        phase = 1.0 * b / abs_b * a.conjugate() / abs_a
-
-    if left:
-        return np.array([[phase * cosine, -sine], [phase * sine, cosine]])
-
-    return np.array([[phase * sine, cosine], [-phase * cosine, sine]])
+    abs_a, abs_b = jnp.abs(a), jnp.abs(b)
+    hypot = jnp.hypot(abs_a, abs_b)
+
+    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)), 
+                    abs_b / hypot, 
+                    jnp.where(jnp.less(abs_a, tol), 1.0, 0.0))
+    
+    sine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+                    abs_a / hypot, 
+                    jnp.where(jnp.less(abs_b, tol), 1.0, 0.0))
+    
+    phase = jnp.where(jnp.logical_and(abs_a > tol, abs_b > tol), 
+                    1.0 * b / abs_b * a.conjugate() / abs_a, 
+                    1.0)
+
+    L = jnp.array([[phase * cosine, -sine], [phase * sine, cosine]])
+    R = jnp.array([[phase * sine, cosine], [-phase * cosine, sine]])
+    return jnp.where(left, L, R)
 
 
 def givens_decomposition(unitary):
@@ -148,7 +151,7 @@ def givens_decomposition(U):
 
     """
 
-    unitary, (M, N) = qml.math.toarray(unitary).copy(), unitary.shape
+    unitary, (M, N) = jnp.copy(unitary), unitary.shape
     if M != N:
         raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")
 
@@ -158,32 +161,33 @@ def givens_decomposition(U):
             for j in range(0, i):
                 indices = [i - j - 1, i - j]
                 grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                unitary[:, indices] = unitary[:, indices] @ grot_mat.T
+                unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
                 right_givens.append((grot_mat.conj(), indices))
         else:
             for j in range(1, i + 1):
                 indices = [N + j - i - 2, N + j - i - 1]
                 grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                unitary[indices] = grot_mat @ unitary[indices]
+                unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
                 left_givens.append((grot_mat, indices))
 
     nleft_givens = []
     for grot_mat, (i, j) in reversed(left_givens):
-        sphase_mat = np.diag(np.diag(unitary)[[i, j]])
+        sphase_mat = jnp.diag(jnp.diag(unitary)[jnp.array([i, j])])
         decomp_mat = grot_mat.conj().T @ sphase_mat
         givens_mat = _givens_matrix(*decomp_mat[1, :].T)
         nphase_mat = decomp_mat @ givens_mat.T
 
         # check for T_{m,n}^{-1} x D = D x T.
-        if not np.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
+        if not jnp.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        unitary[i, i], unitary[j, j] = np.diag(nphase_mat)
+        unitary = unitary.at[i, i].set(jnp.diag(nphase_mat)[0])
+        unitary = unitary.at[j, j].set(jnp.diag(nphase_mat)[1])
         nleft_givens.append((givens_mat.conj(), (i, j)))
 
-    phases, ordered_rotations = np.diag(unitary), []
+    phases, ordered_rotations = jnp.diag(unitary), []
     for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
-        if not np.all(np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])):  # pragma: no cover
+        if not jnp.all(jnp.isreal(grot_mat[0, 1]) and jnp.isreal(grot_mat[1, 1])):  # pragma: no cover
             raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))
 

From c17cf831eb05e7090ef87f60df86bf912c03e0db Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 13:59:01 +0200
Subject: [PATCH 02/55] Converting to jax.numpy the basis rotation

---
 .../templates/subroutines/basis_rotation.py    | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 346e8b8370f..ca0061cd642 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -20,6 +20,7 @@
 import pennylane as qml
 from pennylane.operation import AnyWires, Operation
 from pennylane.qchem.givens_decomposition import givens_decomposition
+import jax.numpy as jnp
 
 
 # pylint: disable-msg=too-many-arguments
@@ -111,14 +112,15 @@ def _primitive_bind_call(cls, wires, unitary_matrix, check=False, id=None):
 
     def __init__(self, wires, unitary_matrix, check=False, id=None):
         M, N = unitary_matrix.shape
+
         if M != N:
             raise ValueError(
                 f"The unitary matrix should be of shape NxN, got {unitary_matrix.shape}"
             )
 
         if check:
-            umat = qml.math.toarray(unitary_matrix)
-            if not np.allclose(umat @ umat.conj().T, np.eye(M, dtype=complex), atol=1e-6):
+            umat = jnp.array(unitary_matrix)
+            if not jnp.allclose(umat @ umat.conj().T, jnp.eye(M, dtype=complex), atol=1e-6):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
@@ -160,8 +162,8 @@ def compute_decomposition(
             )
 
         if check:
-            umat = qml.math.toarray(unitary_matrix)
-            if not np.allclose(umat @ umat.conj().T, np.eye(M, dtype=complex), atol=1e-4):
+            umat = unitary_matrix #qml.math.toarray(unitary_matrix)
+            if not jnp.allclose(umat @ umat.conj().T, jnp.eye(M, dtype=complex), atol=1e-4):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
@@ -171,17 +173,17 @@ def compute_decomposition(
         phase_list, givens_list = givens_decomposition(unitary_matrix)
 
         for idx, phase in enumerate(phase_list):
-            op_list.append(qml.PhaseShift(np.angle(phase), wires=wires[idx]))
+            op_list.append(qml.PhaseShift(jnp.angle(phase), wires=wires[idx]))
 
         for grot_mat, indices in givens_list:
-            theta = np.arccos(np.real(grot_mat[1, 1]))
-            phi = np.angle(grot_mat[0, 0])
+            theta = jnp.arccos(jnp.real(grot_mat[1, 1]))
+            phi = jnp.angle(grot_mat[0, 0])
 
             op_list.append(
                 qml.SingleExcitation(2 * theta, wires=[wires[indices[0]], wires[indices[1]]])
             )
 
-            if not np.isclose(phi, 0.0):
+            if not jnp.isclose(phi, 0.0):
                 op_list.append(qml.PhaseShift(phi, wires=wires[indices[0]]))
 
         return op_list

From cb663d354b71fabe22852485352e541ca5bd8d2a Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 14:03:02 +0200
Subject: [PATCH 03/55] Update basis_rotation.py

---
 pennylane/templates/subroutines/basis_rotation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index ca0061cd642..ae2cf42375e 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -119,7 +119,7 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
             )
 
         if check:
-            umat = jnp.array(unitary_matrix)
+            umat = unitary_matrix
             if not jnp.allclose(umat @ umat.conj().T, jnp.eye(M, dtype=complex), atol=1e-6):
                 raise ValueError("The provided transformation matrix should be unitary.")
 

From cb4d1f73941048f7f42811a850cf6f71d0c029d5 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 14:08:44 +0200
Subject: [PATCH 04/55] Update changelog-dev.md

---
 doc/releases/changelog-dev.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8559de76b3b..523af297d99 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -135,8 +135,8 @@
   [(#5978)](https://github.com/PennyLaneAI/pennylane/pull/5978)
 
 * `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
-[(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
-
+  [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
+* `qml.BasisRotation` is now qjit compatible (#6004)
 
 <h3>Contributors ✍️</h3>
 

From 11a4b2796520dd29f5d7001154409d536db63c39 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 15:12:05 +0200
Subject: [PATCH 05/55] Update givens_decomposition.py

---
 pennylane/qchem/givens_decomposition.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index ca7f6594e4f..1e76ea8e66e 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -44,16 +44,16 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     abs_a, abs_b = jnp.abs(a), jnp.abs(b)
     hypot = jnp.hypot(abs_a, abs_b)
 
-    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)), 
-                    abs_b / hypot, 
+    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+                    abs_b / hypot,
                     jnp.where(jnp.less(abs_a, tol), 1.0, 0.0))
     
     sine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
-                    abs_a / hypot, 
+                    abs_a / hypot,
                     jnp.where(jnp.less(abs_b, tol), 1.0, 0.0))
     
-    phase = jnp.where(jnp.logical_and(abs_a > tol, abs_b > tol), 
-                    1.0 * b / abs_b * a.conjugate() / abs_a, 
+    phase = jnp.where(jnp.logical_and(abs_a > tol, abs_b > tol),
+                    1.0 * b / abs_b * a.conjugate() / abs_a,
                     1.0)
 
     L = jnp.array([[phase * cosine, -sine], [phase * sine, cosine]])

From d811833b0bf55ce5c9dc2351678dba8aa4f9354a Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 16:05:00 +0200
Subject: [PATCH 06/55] Reverting from jnp. to qml.math.

---
 pennylane/qchem/givens_decomposition.py       | 35 +++++++++----------
 .../templates/subroutines/basis_rotation.py   | 16 ++++-----
 2 files changed, 24 insertions(+), 27 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 1e76ea8e66e..67894a0ff60 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -16,7 +16,6 @@
 """
 
 import numpy as np
-import jax.numpy as jnp
 import pennylane as qml
 
 
@@ -41,24 +40,24 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         np.ndarray (or tensor): Givens rotation matrix
 
     """
-    abs_a, abs_b = jnp.abs(a), jnp.abs(b)
-    hypot = jnp.hypot(abs_a, abs_b)
+    abs_a, abs_b = qml.math.abs(a), qml.math.abs(b)
+    hypot = qml.math.hypot(abs_a, abs_b)
 
-    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+    cosine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     abs_b / hypot,
-                    jnp.where(jnp.less(abs_a, tol), 1.0, 0.0))
+                    qml.math.where(qml.math.less(abs_a, tol), 1.0, 0.0))
     
-    sine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+    sine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     abs_a / hypot,
-                    jnp.where(jnp.less(abs_b, tol), 1.0, 0.0))
+                    qml.math.where(qml.math.less(abs_b, tol), 1.0, 0.0))
     
-    phase = jnp.where(jnp.logical_and(abs_a > tol, abs_b > tol),
+    phase = qml.math.where(qml.math.logical_and(abs_a > tol, abs_b > tol),
                     1.0 * b / abs_b * a.conjugate() / abs_a,
                     1.0)
 
-    L = jnp.array([[phase * cosine, -sine], [phase * sine, cosine]])
-    R = jnp.array([[phase * sine, cosine], [-phase * cosine, sine]])
-    return jnp.where(left, L, R)
+    L = qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]])
+    R = qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]])
+    return qml.math.where(left, L, R)
 
 
 def givens_decomposition(unitary):
@@ -151,7 +150,7 @@ def givens_decomposition(U):
 
     """
 
-    unitary, (M, N) = jnp.copy(unitary), unitary.shape
+    unitary, (M, N) = qml.math.copy(unitary), unitary.shape
     if M != N:
         raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")
 
@@ -172,22 +171,22 @@ def givens_decomposition(U):
 
     nleft_givens = []
     for grot_mat, (i, j) in reversed(left_givens):
-        sphase_mat = jnp.diag(jnp.diag(unitary)[jnp.array([i, j])])
+        sphase_mat = qml.math.diag(qml.math.diag(unitary)[qml.math.array([i, j])])
         decomp_mat = grot_mat.conj().T @ sphase_mat
         givens_mat = _givens_matrix(*decomp_mat[1, :].T)
         nphase_mat = decomp_mat @ givens_mat.T
 
         # check for T_{m,n}^{-1} x D = D x T.
-        if not jnp.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
+        if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        unitary = unitary.at[i, i].set(jnp.diag(nphase_mat)[0])
-        unitary = unitary.at[j, j].set(jnp.diag(nphase_mat)[1])
+        unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
+        unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
         nleft_givens.append((givens_mat.conj(), (i, j)))
 
-    phases, ordered_rotations = jnp.diag(unitary), []
+    phases, ordered_rotations = qml.math.diag(unitary), []
     for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
-        if not jnp.all(jnp.isreal(grot_mat[0, 1]) and jnp.isreal(grot_mat[1, 1])):  # pragma: no cover
+        if not qml.math.all(qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])):  # pragma: no cover
             raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))
 
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index ae2cf42375e..58ca9a62ed7 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -20,8 +20,6 @@
 import pennylane as qml
 from pennylane.operation import AnyWires, Operation
 from pennylane.qchem.givens_decomposition import givens_decomposition
-import jax.numpy as jnp
-
 
 # pylint: disable-msg=too-many-arguments
 class BasisRotation(Operation):
@@ -120,7 +118,7 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
 
         if check:
             umat = unitary_matrix
-            if not jnp.allclose(umat @ umat.conj().T, jnp.eye(M, dtype=complex), atol=1e-6):
+            if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-6):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
@@ -162,8 +160,8 @@ def compute_decomposition(
             )
 
         if check:
-            umat = unitary_matrix #qml.math.toarray(unitary_matrix)
-            if not jnp.allclose(umat @ umat.conj().T, jnp.eye(M, dtype=complex), atol=1e-4):
+            umat = unitary_matrix
+            if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
@@ -173,17 +171,17 @@ def compute_decomposition(
         phase_list, givens_list = givens_decomposition(unitary_matrix)
 
         for idx, phase in enumerate(phase_list):
-            op_list.append(qml.PhaseShift(jnp.angle(phase), wires=wires[idx]))
+            op_list.append(qml.PhaseShift(qml.math.angle(phase), wires=wires[idx]))
 
         for grot_mat, indices in givens_list:
-            theta = jnp.arccos(jnp.real(grot_mat[1, 1]))
-            phi = jnp.angle(grot_mat[0, 0])
+            theta = qml.math.arccos(qml.math.real(grot_mat[1, 1]))
+            phi = qml.math.angle(grot_mat[0, 0])
 
             op_list.append(
                 qml.SingleExcitation(2 * theta, wires=[wires[indices[0]], wires[indices[1]]])
             )
 
-            if not jnp.isclose(phi, 0.0):
+            if not qml.math.isclose(phi, 0.0):
                 op_list.append(qml.PhaseShift(phi, wires=wires[indices[0]]))
 
         return op_list

From 7edc75ac01afe342976d60ea7aaeb6ef98b83165 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 16:08:48 +0200
Subject: [PATCH 07/55] Remove import numpy

---
 pennylane/qchem/givens_decomposition.py           | 1 -
 pennylane/templates/subroutines/basis_rotation.py | 2 --
 2 files changed, 3 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 67894a0ff60..82dd40f33a0 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -15,7 +15,6 @@
 This module contains the functions needed for performing basis transformations defined by a set of fermionic ladder operators.
 """
 
-import numpy as np
 import pennylane as qml
 
 
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 58ca9a62ed7..d23aef262f9 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -15,8 +15,6 @@
 This module contains the template for performing basis transformation defined by a set of fermionic ladder operators.
 """
 
-import numpy as np
-
 import pennylane as qml
 from pennylane.operation import AnyWires, Operation
 from pennylane.qchem.givens_decomposition import givens_decomposition

From d47d601fd50c64b9a7cce9460ce9172b65402f9d Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 16:16:02 +0200
Subject: [PATCH 08/55] Unit test for jax jitting

---
 .../test_subroutines/test_basis_rotation.py   | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index c5c9c643450..54f8663c099 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -446,6 +446,51 @@ def test_jax(self, tol):
 
         assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
 
+    @pytest.mark.jax
+    def test_jax_jit(self, tol):
+        """Test the jax interface."""
+
+        import jax
+        import jax.numpy as jnp
+
+        unitary_matrix = jnp.array(
+            [
+                [0.51378719 + 0.0j, 0.0546265 + 0.79145487j, -0.2051466 + 0.2540723j],
+                [0.62651582 + 0.0j, -0.00828925 - 0.60570321j, -0.36704948 + 0.32528067j],
+                [-0.58608928 + 0.0j, 0.03902657 + 0.04633548j, -0.57220635 + 0.57044649j],
+            ]
+        )
+        weights = jnp.array(
+            [
+                2.2707802713289267,
+                2.9355948424220206,
+                -1.4869222527726533,
+                1.2601662579297865,
+                2.3559705032936717,
+                1.1748572730890159,
+                2.2500537657656356,
+                -0.7251404204443089,
+                2.3577346350335198,
+            ]
+        )
+
+        dev = qml.device("default.qubit", wires=3)
+
+        circuit = qml.QNode(circuit_template, dev)
+        circuit2 = qml.QNode(circuit_decomposed, dev)
+
+        res = jax.jit(circuit)(unitary_matrix)
+        res2 = jax.jit(circuit2)(weights)
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+        grad_fn = jax.grad(circuit)
+        grads = grad_fn(unitary_matrix)
+
+        grad_fn2 = jax.grad(circuit2)
+        grads2 = grad_fn2(weights)
+
+        assert qml.math.allclose(grads[0], grads2[0], atol=tol, rtol=0)
+
     @pytest.mark.tf
     def test_tf(self, tol):
         """Test the tf interface."""

From ca95f7a7e5b745fa7de145102e39f2957a862549 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 16:20:20 +0200
Subject: [PATCH 09/55] Removing whitespaces

---
 pennylane/qchem/givens_decomposition.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 82dd40f33a0..92b5204f0c3 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -45,11 +45,11 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     cosine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     abs_b / hypot,
                     qml.math.where(qml.math.less(abs_a, tol), 1.0, 0.0))
-    
+
     sine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     abs_a / hypot,
                     qml.math.where(qml.math.less(abs_b, tol), 1.0, 0.0))
-    
+
     phase = qml.math.where(qml.math.logical_and(abs_a > tol, abs_b > tol),
                     1.0 * b / abs_b * a.conjugate() / abs_a,
                     1.0)

From eb77772841208717dc8e4d9d66bdff5fc6ddf8c6 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 16:28:24 +0200
Subject: [PATCH 10/55] qml.math.copy and (not
 qml.math.is_abstract(unitary_matrix))

---
 pennylane/qchem/givens_decomposition.py           |  4 ++--
 pennylane/templates/subroutines/basis_rotation.py | 10 +++++-----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 92b5204f0c3..60b285d753e 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -149,9 +149,9 @@ def givens_decomposition(U):
 
     """
 
-    unitary, (M, N) = qml.math.copy(unitary), unitary.shape
+    unitary, (M, N) = qml.math.copy(unitary), qml.math.shape(unitary)
     if M != N:
-        raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")
+        raise ValueError(f"The unitary matrix should be of shape NxN, got {qml.math.shape(unitary)}")
 
     left_givens, right_givens = [], []
     for i in range(1, N):
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index d23aef262f9..f6da354b4e1 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -107,7 +107,7 @@ def _primitive_bind_call(cls, wires, unitary_matrix, check=False, id=None):
         return cls._primitive.bind(*wires, unitary_matrix, check=check, id=id)
 
     def __init__(self, wires, unitary_matrix, check=False, id=None):
-        M, N = unitary_matrix.shape
+        M, N = qml.math.shape(unitary_matrix)
 
         if M != N:
             raise ValueError(
@@ -115,7 +115,7 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
             )
 
         if check:
-            umat = unitary_matrix
+            umat = qml.math.copy(unitary_matrix)
             if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-6):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
@@ -151,15 +151,15 @@ def compute_decomposition(
             list[.Operator]: decomposition of the operator
         """
 
-        M, N = unitary_matrix.shape
+        M, N = qml.math.shape(unitary_matrix)
         if M != N:
             raise ValueError(
                 f"The unitary matrix should be of shape NxN, got {unitary_matrix.shape}"
             )
 
         if check:
-            umat = unitary_matrix
-            if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4):
+            umat = qml.math.copy(unitary_matrix)
+            if (not qml.math.is_abstract(unitary_matrix)) and (not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4)):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:

From 7d6c223d59e37ede3d091357621072bc3959abc4 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 16:31:31 +0200
Subject: [PATCH 11/55] Update basis_rotation.py

---
 pennylane/templates/subroutines/basis_rotation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index f6da354b4e1..9b4e6a8c354 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -111,7 +111,7 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
 
         if M != N:
             raise ValueError(
-                f"The unitary matrix should be of shape NxN, got {unitary_matrix.shape}"
+                f"The unitary matrix should be of shape NxN, got {qml.math.shape(unitary_matrix)}"
             )
 
         if check:

From 38e3e68da70a06dea5fec9de30cef01de5933831 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 16:53:43 +0200
Subject: [PATCH 12/55] branching due to .at

---
 pennylane/qchem/givens_decomposition.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 60b285d753e..a1d36fa61ca 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -159,13 +159,19 @@ def givens_decomposition(U):
             for j in range(0, i):
                 indices = [i - j - 1, i - j]
                 grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
+                if qml.math.get_interface(unitary) == 'jax':
+                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
+                else:
+                    unitary[indices, :] = grot_mat @ unitary[indices, :]
                 right_givens.append((grot_mat.conj(), indices))
         else:
             for j in range(1, i + 1):
                 indices = [N + j - i - 2, N + j - i - 1]
                 grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
+                if qml.math.get_interface(unitary) == 'jax':
+                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
+                else:
+                    unitary[indices, :] = grot_mat @ unitary[indices, :]
                 left_givens.append((grot_mat, indices))
 
     nleft_givens = []
@@ -179,8 +185,11 @@ def givens_decomposition(U):
         if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
-        unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
+        if qml.math.get_interface(unitary) == 'jax':
+            unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
+            unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
+        else:
+            unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
         nleft_givens.append((givens_mat.conj(), (i, j)))
 
     phases, ordered_rotations = qml.math.diag(unitary), []

From 07816dffeaa1657d67db3e77d70e16521802e852 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 17:23:34 +0200
Subject: [PATCH 13/55] Branching functions with/without jax backend

---
 pennylane/qchem/givens_decomposition.py       | 190 ++++++++++++++++--
 .../templates/subroutines/basis_rotation.py   |   7 +-
 2 files changed, 182 insertions(+), 15 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index a1d36fa61ca..8d7083c9dab 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -16,6 +16,7 @@
 """
 
 import pennylane as qml
+import jax.numpy as jnp
 
 
 def _givens_matrix(a, b, left=True, tol=1e-8):
@@ -50,7 +51,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
                     abs_a / hypot,
                     qml.math.where(qml.math.less(abs_b, tol), 1.0, 0.0))
 
-    phase = qml.math.where(qml.math.logical_and(abs_a > tol, abs_b > tol),
+    phase = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     1.0 * b / abs_b * a.conjugate() / abs_a,
                     1.0)
 
@@ -58,6 +59,45 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     R = qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]])
     return qml.math.where(left, L, R)
 
+def _givens_matrix_jax(a, b, left=True, tol=1e-8):
+    r"""Build a :math:`2 \times 2` Givens rotation matrix :math:`G`.
+
+    When the matrix :math:`G` is applied to a vector :math:`[a,\ b]^T` the following would happen:
+
+    .. math::
+
+            G \times \begin{bmatrix} a \\ b \end{bmatrix} = \begin{bmatrix} 0 \\ r \end{bmatrix} \quad \quad \quad \begin{bmatrix} a \\ b \end{bmatrix} \times G = \begin{bmatrix} r \\ 0 \end{bmatrix},
+
+    where :math:`r` is a complex number.
+
+    Args:
+        a (float or complex): first element of the vector for which the Givens matrix is being computed
+        b (float or complex): second element of the vector for which the Givens matrix is being computed
+        left (bool): determines if the Givens matrix is being applied from the left side or right side.
+        tol (float): determines tolerance limits for :math:`|a|` and :math:`|b|` under which they are considered as zero.
+
+    Returns:
+        np.ndarray (or tensor): Givens rotation matrix
+
+    """
+    abs_a, abs_b = jnp.abs(a), jnp.abs(b)
+    hypot = jnp.hypot(abs_a, abs_b)
+
+    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+                    abs_b / hypot,
+                    jnp.where(jnp.less(abs_a, tol), 1.0, 0.0))
+
+    sine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+                    abs_a / hypot,
+                    jnp.where(jnp.less(abs_b, tol), 1.0, 0.0))
+
+    phase = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+                    1.0 * b / abs_b * a.conjugate() / abs_a,
+                    1.0)
+
+    L = jnp.array([[phase * cosine, -sine], [phase * sine, cosine]])
+    R = jnp.array([[phase * sine, cosine], [-phase * cosine, sine]])
+    return jnp.where(left, L, R)
 
 def givens_decomposition(unitary):
     r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.
@@ -159,19 +199,13 @@ def givens_decomposition(U):
             for j in range(0, i):
                 indices = [i - j - 1, i - j]
                 grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                if qml.math.get_interface(unitary) == 'jax':
-                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
-                else:
-                    unitary[indices, :] = grot_mat @ unitary[indices, :]
+                unitary[indices, :] = grot_mat @ unitary[indices, :]
                 right_givens.append((grot_mat.conj(), indices))
         else:
             for j in range(1, i + 1):
                 indices = [N + j - i - 2, N + j - i - 1]
                 grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                if qml.math.get_interface(unitary) == 'jax':
-                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
-                else:
-                    unitary[indices, :] = grot_mat @ unitary[indices, :]
+                unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
                 left_givens.append((grot_mat, indices))
 
     nleft_givens = []
@@ -185,11 +219,141 @@ def givens_decomposition(U):
         if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        if qml.math.get_interface(unitary) == 'jax':
-            unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
-            unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
+        unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
+        nleft_givens.append((givens_mat.conj(), (i, j)))
+
+    phases, ordered_rotations = qml.math.diag(unitary), []
+    for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
+        if not qml.math.all(qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])):  # pragma: no cover
+            raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
+        ordered_rotations.append((grot_mat, (i, j)))
+
+    return phases, ordered_rotations
+
+def givens_decomposition_jax(unitary):
+    r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.
+
+    This decomposition is based on the construction scheme given in `Optica, 3, 1460 (2016) <https://opg.optica.org/optica/fulltext.cfm?uri=optica-3-12-1460&id=355743>`_\ ,
+    which allows one to write any unitary matrix :math:`U` as:
+
+    .. math::
+
+        U = D \left(\prod_{(m, n) \in G} T_{m, n}(\theta, \phi)\right),
+
+    where :math:`D` is a diagonal phase matrix, :math:`T(\theta, \phi)` is the Givens rotation gates with phase shifts and :math:`G` defines the
+    specific ordered sequence of the Givens rotation gates acting on wires :math:`(m, n)`. The unitary for the :math:`T(\theta, \phi)` gates
+    appearing in the decomposition is of the following form:
+
+    .. math:: T(\theta, \phi) = \begin{bmatrix}
+                                    1 & 0 & 0 & 0 \\
+                                    0 & e^{i \phi} \cos(\theta) & -\sin(\theta) & 0 \\
+                                    0 & e^{i \phi} \sin(\theta) & \cos(\theta) & 0 \\
+                                    0 & 0 & 0 & 1
+                                \end{bmatrix},
+
+    where :math:`\theta \in [0, \pi/2]` is the angle of rotation in the :math:`\{|01\rangle, |10\rangle \}` subspace
+    and :math:`\phi \in [0, 2 \pi]` represents the phase shift at the first wire.
+
+    This function is to be used with JAX interface.
+
+    **Example**
+
+    .. code-block:: python
+
+        unitary = np.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
+                            [-0.21271+0.34938j, -0.38853+0.36497j,  0.61467-0.41317j],
+                            [ 0.41356-0.20765j, -0.00651-0.66689j,  0.32839-0.48293j]])
+
+        phase_mat, ordered_rotations = givens_decomposition(unitary)
+
+    >>> phase_mat
+    tensor([-0.20604358+0.9785369j , -0.82993272+0.55786114j,
+            0.56230612-0.82692833j], requires_grad=True)
+    >>> ordered_rotations
+    [(tensor([[-0.65087861-0.63937521j, -0.40933651-0.j        ],
+            [-0.29201359-0.28685265j,  0.91238348-0.j        ]], requires_grad=True),
+    (0, 1)),
+    (tensor([[ 0.47970366-0.33308926j, -0.8117487 -0.j        ],
+            [ 0.66677093-0.46298215j,  0.5840069 -0.j        ]], requires_grad=True),
+    (1, 2)),
+    (tensor([[ 0.36147547+0.73779454j, -0.57008306-0.j        ],
+            [ 0.2508207 +0.51194108j,  0.82158706-0.j        ]], requires_grad=True),
+    (0, 1))]
+
+    Args:
+        unitary (tensor): unitary matrix on which decomposition will be performed
+
+    Returns:
+        (np.ndarray, list[(np.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
+
+    Raises:
+        ValueError: if the provided matrix is not square.
+
+    .. details::
+        :title: Theory and Pseudocode
+
+        **Givens Rotation**
+
+        Applying the Givens rotation :math:`T(\theta, \phi)` performs the following transformation of the basis states:
+
+        .. math::
+
+            &|00\rangle \mapsto |00\rangle\\
+            &|01\rangle \mapsto e^{i \phi} \cos(\theta) |01\rangle - \sin(\theta) |10\rangle\\
+            &|10\rangle \mapsto e^{i \phi} \sin(\theta) |01\rangle + \cos(\theta) |10\rangle\\
+            &|11\rangle \mapsto |11\rangle.
+
+        **Pseudocode**
+
+        The algorithm that implements the decomposition is the following:
+
+        .. code-block:: python
+
+            def givens_decomposition(U):
+                for i in range(1, N):
+                    if i % 2:
+                        for j in range(0, i):
+                            # Find T^-1(i-j, i-j+1) matrix that nulls element (N-j, i-j) of U
+                            # Update U = U @ T^-1(i-j, i-j+1)
+                    else:
+                        for j in range(1, i):
+                            # Find T(N+j-i-1, N+j-i) matrix that nulls element (N+j-i, j) of U
+                            # Update U = T(N+j-i-1, N+j-i) @ U
+
+    """
+
+    unitary, (M, N) = qml.math.copy(unitary), qml.math.shape(unitary)
+    if M != N:
+        raise ValueError(f"The unitary matrix should be of shape NxN, got {qml.math.shape(unitary)}")
+
+    left_givens, right_givens = [], []
+    for i in range(1, N):
+        if i % 2:
+            for j in range(0, i):
+                indices = [i - j - 1, i - j]
+                grot_mat = _givens_matrix_jax(*unitary[N - j - 1, indices].T, left=True)
+                unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
+                right_givens.append((grot_mat.conj(), indices))
         else:
-            unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
+            for j in range(1, i + 1):
+                indices = [N + j - i - 2, N + j - i - 1]
+                grot_mat = _givens_matrix_jax(*unitary[indices, j - 1], left=False)
+                unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
+                left_givens.append((grot_mat, indices))
+
+    nleft_givens = []
+    for grot_mat, (i, j) in reversed(left_givens):
+        sphase_mat = qml.math.diag(qml.math.diag(unitary)[qml.math.array([i, j])])
+        decomp_mat = grot_mat.conj().T @ sphase_mat
+        givens_mat = _givens_matrix_jax(*decomp_mat[1, :].T)
+        nphase_mat = decomp_mat @ givens_mat.T
+
+        # check for T_{m,n}^{-1} x D = D x T.
+        if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
+            raise ValueError("Failed to shift phase transposition.")
+
+        unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
+        unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
         nleft_givens.append((givens_mat.conj(), (i, j)))
 
     phases, ordered_rotations = qml.math.diag(unitary), []
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 9b4e6a8c354..3a73ee30c72 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -17,7 +17,7 @@
 
 import pennylane as qml
 from pennylane.operation import AnyWires, Operation
-from pennylane.qchem.givens_decomposition import givens_decomposition
+from pennylane.qchem.givens_decomposition import givens_decomposition, givens_decomposition_jax
 
 # pylint: disable-msg=too-many-arguments
 class BasisRotation(Operation):
@@ -166,7 +166,10 @@ def compute_decomposition(
             raise ValueError(f"This template requires at least two wires, got {len(wires)}")
 
         op_list = []
-        phase_list, givens_list = givens_decomposition(unitary_matrix)
+        if qml.math.get_interface(unitary_matrix) == 'jax':
+            phase_list, givens_list = givens_decomposition_jax(unitary_matrix)
+        else:
+            phase_list, givens_list = givens_decomposition(unitary_matrix)
 
         for idx, phase in enumerate(phase_list):
             op_list.append(qml.PhaseShift(qml.math.angle(phase), wires=wires[idx]))

From cd76203ceef8090e4f4263f9eedd22615428b18f Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 17:34:48 +0200
Subject: [PATCH 14/55] Update givens_decomposition.py

---
 pennylane/qchem/givens_decomposition.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 8d7083c9dab..1765ab1fce9 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -16,7 +16,6 @@
 """
 
 import pennylane as qml
-import jax.numpy as jnp
 
 
 def _givens_matrix(a, b, left=True, tol=1e-8):
@@ -80,24 +79,24 @@ def _givens_matrix_jax(a, b, left=True, tol=1e-8):
         np.ndarray (or tensor): Givens rotation matrix
 
     """
-    abs_a, abs_b = jnp.abs(a), jnp.abs(b)
-    hypot = jnp.hypot(abs_a, abs_b)
+    abs_a, abs_b = qml.math.abs(a), qml.math.abs(b)
+    hypot = qml.math.hypot(abs_a, abs_b)
 
-    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+    cosine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     abs_b / hypot,
-                    jnp.where(jnp.less(abs_a, tol), 1.0, 0.0))
+                    qml.math.where(qml.math.less(abs_a, tol), 1.0, 0.0))
 
-    sine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+    sine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     abs_a / hypot,
-                    jnp.where(jnp.less(abs_b, tol), 1.0, 0.0))
+                    qml.math.where(qml.math.less(abs_b, tol), 1.0, 0.0))
 
-    phase = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
+    phase = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
                     1.0 * b / abs_b * a.conjugate() / abs_a,
                     1.0)
 
-    L = jnp.array([[phase * cosine, -sine], [phase * sine, cosine]])
-    R = jnp.array([[phase * sine, cosine], [-phase * cosine, sine]])
-    return jnp.where(left, L, R)
+    L = qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]])
+    R = qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]])
+    return qml.math.where(left, L, R)
 
 def givens_decomposition(unitary):
     r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.

From 75e4f1e73b575b25e285f55ddcde60fa5800de14 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 17:56:24 +0200
Subject: [PATCH 15/55] Update test_basis_rotation.py

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

diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index 54f8663c099..d0cd157c06c 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -335,7 +335,7 @@ def test_id(self):
 
 
 def circuit_template(unitary_matrix):
-    qml.BasisState(np.array([1, 1, 0]), wires=[0, 1, 2])
+    qml.BasisState(qml.math.array([1, 1, 0]), wires=[0, 1, 2])
     qml.BasisRotation(
         wires=range(3),
         unitary_matrix=unitary_matrix,

From 8bfba222026fb2cee9723a91559edc3f10150322 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 18:20:28 +0200
Subject: [PATCH 16/55] Update givens_decomposition.py

---
 pennylane/qchem/givens_decomposition.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 1765ab1fce9..6067da38a86 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -79,24 +79,25 @@ def _givens_matrix_jax(a, b, left=True, tol=1e-8):
         np.ndarray (or tensor): Givens rotation matrix
 
     """
-    abs_a, abs_b = qml.math.abs(a), qml.math.abs(b)
-    hypot = qml.math.hypot(abs_a, abs_b)
+    import jax.numpy as jnp
+    abs_a, abs_b = jnp.abs(a), jnp.abs(b)
+    hypot = jnp.hypot(abs_a, abs_b)
 
-    cosine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
+    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
                     abs_b / hypot,
-                    qml.math.where(qml.math.less(abs_a, tol), 1.0, 0.0))
+                    jnp.where(jnp.less(abs_a, tol), 1.0, 0.0))
 
-    sine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
+    sine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
                     abs_a / hypot,
-                    qml.math.where(qml.math.less(abs_b, tol), 1.0, 0.0))
+                    jnp.where(jnp.less(abs_b, tol), 1.0, 0.0))
 
-    phase = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
+    phase = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
                     1.0 * b / abs_b * a.conjugate() / abs_a,
                     1.0)
 
-    L = qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]])
-    R = qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]])
-    return qml.math.where(left, L, R)
+    L = jnp.array([[phase * cosine, -sine], [phase * sine, cosine]])
+    R = jnp.array([[phase * sine, cosine], [-phase * cosine, sine]])
+    return jnp.where(left, L, R)
 
 def givens_decomposition(unitary):
     r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.

From 34d42578636c9ee7349160180e44cce181352a61 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 18:42:06 +0200
Subject: [PATCH 17/55] Update givens_decomposition.py

---
 pennylane/qchem/givens_decomposition.py | 53 ++++++++++++-------------
 1 file changed, 26 insertions(+), 27 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 6067da38a86..54602809cf2 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -16,6 +16,7 @@
 """
 
 import pennylane as qml
+import numpy as np
 
 
 def _givens_matrix(a, b, left=True, tol=1e-8):
@@ -39,24 +40,21 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         np.ndarray (or tensor): Givens rotation matrix
 
     """
-    abs_a, abs_b = qml.math.abs(a), qml.math.abs(b)
-    hypot = qml.math.hypot(abs_a, abs_b)
-
-    cosine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
-                    abs_b / hypot,
-                    qml.math.where(qml.math.less(abs_a, tol), 1.0, 0.0))
-
-    sine = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
-                    abs_a / hypot,
-                    qml.math.where(qml.math.less(abs_b, tol), 1.0, 0.0))
-
-    phase = qml.math.where(qml.math.logical_and(qml.math.greater(abs_a, tol), qml.math.greater(abs_b, tol)),
-                    1.0 * b / abs_b * a.conjugate() / abs_a,
-                    1.0)
-
-    L = qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]])
-    R = qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]])
-    return qml.math.where(left, L, R)
+    abs_a, abs_b = np.abs(a), np.abs(b)
+    if abs_a < tol:
+        cosine, sine, phase = 1.0, 0.0, 1.0
+    elif abs_b < tol:
+        cosine, sine, phase = 0.0, 1.0, 1.0
+    else:
+        hypot = np.hypot(abs_a, abs_b)
+        cosine = abs_b / hypot
+        sine = abs_a / hypot
+        phase = 1.0 * b / abs_b * a.conjugate() / abs_a
+
+    if left:
+        return np.array([[phase * cosine, -sine], [phase * sine, cosine]])
+
+    return np.array([[phase * sine, cosine], [-phase * cosine, sine]])
 
 def _givens_matrix_jax(a, b, left=True, tol=1e-8):
     r"""Build a :math:`2 \times 2` Givens rotation matrix :math:`G`.
@@ -189,9 +187,9 @@ def givens_decomposition(U):
 
     """
 
-    unitary, (M, N) = qml.math.copy(unitary), qml.math.shape(unitary)
+    unitary, (M, N) = qml.math.toarray(unitary).copy(), unitary.shape
     if M != N:
-        raise ValueError(f"The unitary matrix should be of shape NxN, got {qml.math.shape(unitary)}")
+        raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")
 
     left_givens, right_givens = [], []
     for i in range(1, N):
@@ -199,37 +197,38 @@ def givens_decomposition(U):
             for j in range(0, i):
                 indices = [i - j - 1, i - j]
                 grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                unitary[indices, :] = grot_mat @ unitary[indices, :]
+                unitary[:, indices] = unitary[:, indices] @ grot_mat.T
                 right_givens.append((grot_mat.conj(), indices))
         else:
             for j in range(1, i + 1):
                 indices = [N + j - i - 2, N + j - i - 1]
                 grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
+                unitary[indices] = grot_mat @ unitary[indices]
                 left_givens.append((grot_mat, indices))
 
     nleft_givens = []
     for grot_mat, (i, j) in reversed(left_givens):
-        sphase_mat = qml.math.diag(qml.math.diag(unitary)[qml.math.array([i, j])])
+        sphase_mat = np.diag(np.diag(unitary)[[i, j]])
         decomp_mat = grot_mat.conj().T @ sphase_mat
         givens_mat = _givens_matrix(*decomp_mat[1, :].T)
         nphase_mat = decomp_mat @ givens_mat.T
 
         # check for T_{m,n}^{-1} x D = D x T.
-        if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
+        if not np.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
+        unitary[i, i], unitary[j, j] = np.diag(nphase_mat)
         nleft_givens.append((givens_mat.conj(), (i, j)))
 
-    phases, ordered_rotations = qml.math.diag(unitary), []
+    phases, ordered_rotations = np.diag(unitary), []
     for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
-        if not qml.math.all(qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])):  # pragma: no cover
+        if not np.all(np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])):  # pragma: no cover
             raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))
 
     return phases, ordered_rotations
 
+
 def givens_decomposition_jax(unitary):
     r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.
 

From 96028a2a0cf8ba3531978d137645b175e5caa295 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 19:10:28 +0200
Subject: [PATCH 18/55] Update test_basis_rotation.py

---
 .../templates/test_subroutines/test_basis_rotation.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index d0cd157c06c..6e9866ad2be 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -334,11 +334,12 @@ def test_id(self):
         assert template.id == "a"
 
 
-def circuit_template(unitary_matrix):
+def circuit_template(unitary_matrix, check = False):
     qml.BasisState(qml.math.array([1, 1, 0]), wires=[0, 1, 2])
     qml.BasisRotation(
         wires=range(3),
         unitary_matrix=unitary_matrix,
+        check = check,
     )
     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
@@ -476,12 +477,12 @@ def test_jax_jit(self, tol):
 
         dev = qml.device("default.qubit", wires=3)
 
-        circuit = qml.QNode(circuit_template, dev)
+        circuit = qml.QNode(jax.jit(circuit_template, static_argnames='check'), dev)
         circuit2 = qml.QNode(circuit_decomposed, dev)
 
-        res = jax.jit(circuit)(unitary_matrix)
-        res2 = jax.jit(circuit2)(weights)
-        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+        res = circuit(unitary_matrix)
+        res2 = circuit2(weights)
+        assert jnp.allclose(res, res2, atol=tol, rtol=0)
 
         grad_fn = jax.grad(circuit)
         grads = grad_fn(unitary_matrix)

From bc34c6f2ae7ed2386486d7d2ebd0a187dcf2bfdb Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 19:10:30 +0200
Subject: [PATCH 19/55] Update basis_rotation.py

---
 pennylane/templates/subroutines/basis_rotation.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 3a73ee30c72..f04b7f86bb8 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -159,17 +159,19 @@ def compute_decomposition(
 
         if check:
             umat = qml.math.copy(unitary_matrix)
-            if (not qml.math.is_abstract(unitary_matrix)) and (not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4)):
-                raise ValueError("The provided transformation matrix should be unitary.")
+            if (not qml.math.is_abstract(unitary_matrix)):
+                if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4):
+                    raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
             raise ValueError(f"This template requires at least two wires, got {len(wires)}")
 
         op_list = []
-        if qml.math.get_interface(unitary_matrix) == 'jax':
-            phase_list, givens_list = givens_decomposition_jax(unitary_matrix)
-        else:
-            phase_list, givens_list = givens_decomposition(unitary_matrix)
+
+        phase_list, givens_list = qml.math.where(
+                                    qml.math.get_interface(unitary_matrix) == 'jax',
+                                    givens_decomposition_jax(unitary_matrix),
+                                    givens_decomposition(unitary_matrix))
 
         for idx, phase in enumerate(phase_list):
             op_list.append(qml.PhaseShift(qml.math.angle(phase), wires=wires[idx]))

From 57bfc8ade51787e98bb1c48c12ba8035c16806bd Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 19:13:43 +0200
Subject: [PATCH 20/55] Update changelog-dev.md

---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 41c5413211c..4429db460af 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -138,7 +138,7 @@
 
 * `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
   [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
-* `qml.BasisRotation` is now qjit compatible (#6004)
+* `qml.BasisRotation` is now qjit compatible ([#6004](`qml.BasisRotation` is now jit compatible ([#6004](https://github.com/PennyLaneAI/pennylane/issues/6004))))
 
 <h3>Contributors ✍️</h3>
 

From 650379c52b29c1475373d607eb6405230dd5ab54 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 19:15:33 +0200
Subject: [PATCH 21/55] Update changelog-dev.md

---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4429db460af..4bd9a7c5931 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -138,7 +138,7 @@
 
 * `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
   [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
-* `qml.BasisRotation` is now qjit compatible ([#6004](`qml.BasisRotation` is now jit compatible ([#6004](https://github.com/PennyLaneAI/pennylane/issues/6004))))
+* `qml.BasisRotation` is now qjit compatible. ([#6004](`qml.BasisRotation` is now jit compatible ([#6004](https://github.com/PennyLaneAI/pennylane/issues/6004))))
 
 <h3>Contributors ✍️</h3>
 

From a5637542a8e7de3db93c46da60da5caf00cf6764 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 22 Jul 2024 19:21:17 +0200
Subject: [PATCH 22/55] Checks for abstract arrays

---
 pennylane/qchem/givens_decomposition.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 54602809cf2..9dc5758e341 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -348,17 +348,19 @@ def givens_decomposition(U):
         nphase_mat = decomp_mat @ givens_mat.T
 
         # check for T_{m,n}^{-1} x D = D x T.
-        if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
-            raise ValueError("Failed to shift phase transposition.")
+        if not qml.math.is_abstract(nphase_mat @ givens_mat.conj()) and not qml.math.is_abstract(decomp_mat):
+            if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
+                raise ValueError("Failed to shift phase transposition.")
 
         unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
         unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
         nleft_givens.append((givens_mat.conj(), (i, j)))
 
     phases, ordered_rotations = qml.math.diag(unitary), []
-    for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
-        if not qml.math.all(qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])):  # pragma: no cover
-            raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
+    if not qml.math.is_abstract(grot_mat):
+        for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
+            if not qml.math.all(qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])):  # pragma: no cover
+                raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))
 
     return phases, ordered_rotations

From e33bf03d41826dbd553324cd27e0d3288e9ecc96 Mon Sep 17 00:00:00 2001
From: obliviateandsurrender <utkarshazad98@gmail.com>
Date: Tue, 23 Jul 2024 01:26:32 -0400
Subject: [PATCH 23/55] make jax-jit friendly

---
 pennylane/qchem/givens_decomposition.py       | 267 ++++--------------
 .../templates/subroutines/basis_rotation.py   |  20 +-
 .../test_subroutines/test_basis_rotation.py   |  70 +----
 3 files changed, 75 insertions(+), 282 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 9dc5758e341..8770c469ec2 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -14,9 +14,9 @@
 """
 This module contains the functions needed for performing basis transformations defined by a set of fermionic ladder operators.
 """
+import numpy as np
 
 import pennylane as qml
-import numpy as np
 
 
 def _givens_matrix(a, b, left=True, tol=1e-8):
@@ -40,63 +40,24 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         np.ndarray (or tensor): Givens rotation matrix
 
     """
-    abs_a, abs_b = np.abs(a), np.abs(b)
-    if abs_a < tol:
-        cosine, sine, phase = 1.0, 0.0, 1.0
-    elif abs_b < tol:
-        cosine, sine, phase = 0.0, 1.0, 1.0
-    else:
-        hypot = np.hypot(abs_a, abs_b)
-        cosine = abs_b / hypot
-        sine = abs_a / hypot
-        phase = 1.0 * b / abs_b * a.conjugate() / abs_a
-
-    if left:
-        return np.array([[phase * cosine, -sine], [phase * sine, cosine]])
-
-    return np.array([[phase * sine, cosine], [-phase * cosine, sine]])
-
-def _givens_matrix_jax(a, b, left=True, tol=1e-8):
-    r"""Build a :math:`2 \times 2` Givens rotation matrix :math:`G`.
-
-    When the matrix :math:`G` is applied to a vector :math:`[a,\ b]^T` the following would happen:
-
-    .. math::
-
-            G \times \begin{bmatrix} a \\ b \end{bmatrix} = \begin{bmatrix} 0 \\ r \end{bmatrix} \quad \quad \quad \begin{bmatrix} a \\ b \end{bmatrix} \times G = \begin{bmatrix} r \\ 0 \end{bmatrix},
-
-    where :math:`r` is a complex number.
-
-    Args:
-        a (float or complex): first element of the vector for which the Givens matrix is being computed
-        b (float or complex): second element of the vector for which the Givens matrix is being computed
-        left (bool): determines if the Givens matrix is being applied from the left side or right side.
-        tol (float): determines tolerance limits for :math:`|a|` and :math:`|b|` under which they are considered as zero.
-
-    Returns:
-        np.ndarray (or tensor): Givens rotation matrix
-
-    """
-    import jax.numpy as jnp
-    abs_a, abs_b = jnp.abs(a), jnp.abs(b)
-    hypot = jnp.hypot(abs_a, abs_b)
+    abs_a, abs_b, interface = qml.math.abs(a), qml.math.abs(b), qml.math.get_interface(a)
+    hypot = qml.math.hypot(abs_a, abs_b)
 
-    cosine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
-                    abs_b / hypot,
-                    jnp.where(jnp.less(abs_a, tol), 1.0, 0.0))
+    cosine = qml.math.where(abs_a < tol, 1.0, qml.math.where(abs_b < tol, 0.0, abs_b / hypot))
+    sine = qml.math.where(abs_a < tol, 0.0, qml.math.where(abs_b < tol, 1.0, abs_a / hypot))
+    phase = qml.math.where(
+        abs_a < tol,
+        1.0,
+        qml.math.where(abs_b < tol, 1.0, 1.0 * b / abs_b * qml.math.conj(a) / abs_a),
+    )
 
-    sine = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
-                    abs_a / hypot,
-                    jnp.where(jnp.less(abs_b, tol), 1.0, 0.0))
+    if left:
+        return qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]], like=interface)
 
-    phase = jnp.where(jnp.logical_and(jnp.greater(abs_a, tol), jnp.greater(abs_b, tol)),
-                    1.0 * b / abs_b * a.conjugate() / abs_a,
-                    1.0)
+    return qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]], like=interface)
 
-    L = jnp.array([[phase * cosine, -sine], [phase * sine, cosine]])
-    R = jnp.array([[phase * sine, cosine], [-phase * cosine, sine]])
-    return jnp.where(left, L, R)
 
+# pylint:disable = too-many-branches
 def givens_decomposition(unitary):
     r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.
 
@@ -186,181 +147,69 @@ def givens_decomposition(U):
                             # Update U = T(N+j-i-1, N+j-i) @ U
 
     """
+    interface = qml.math.get_deep_interface(unitary)
+    unitary = qml.math.copy(unitary) if interface == "jax" else qml.math.toarray(unitary).copy()
+    M, N = qml.math.shape(unitary)
 
-    unitary, (M, N) = qml.math.toarray(unitary).copy(), unitary.shape
     if M != N:
         raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")
 
     left_givens, right_givens = [], []
-    for i in range(1, N):
-        if i % 2:
-            for j in range(0, i):
-                indices = [i - j - 1, i - j]
-                grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                unitary[:, indices] = unitary[:, indices] @ grot_mat.T
-                right_givens.append((grot_mat.conj(), indices))
-        else:
-            for j in range(1, i + 1):
-                indices = [N + j - i - 2, N + j - i - 1]
-                grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                unitary[indices] = grot_mat @ unitary[indices]
-                left_givens.append((grot_mat, indices))
+    if interface == "jax":
+        for i in range(1, N):
+            if i % 2:
+                for j in range(0, i):
+                    indices = [i - j - 1, i - j]
+                    grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
+                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
+                    right_givens.append((grot_mat.conj(), indices))
+            else:
+                for j in range(1, i + 1):
+                    indices = [N + j - i - 2, N + j - i - 1]
+                    grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
+                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
+                    left_givens.append((grot_mat, indices))
+    else:
+        for i in range(1, N):
+            if i % 2:
+                for j in range(0, i):
+                    indices = [i - j - 1, i - j]
+                    grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
+                    unitary[:, indices] = unitary[:, indices] @ grot_mat.T
+                    right_givens.append((grot_mat.conj(), indices))
+            else:
+                for j in range(1, i + 1):
+                    indices = [N + j - i - 2, N + j - i - 1]
+                    grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
+                    unitary[indices] = grot_mat @ unitary[indices]
+                    left_givens.append((grot_mat, indices))
 
     nleft_givens = []
     for grot_mat, (i, j) in reversed(left_givens):
-        sphase_mat = np.diag(np.diag(unitary)[[i, j]])
+        sphase_mat = qml.math.diag(qml.math.diag(unitary)[qml.math.array([i, j])])
         decomp_mat = grot_mat.conj().T @ sphase_mat
         givens_mat = _givens_matrix(*decomp_mat[1, :].T)
         nphase_mat = decomp_mat @ givens_mat.T
 
         # check for T_{m,n}^{-1} x D = D x T.
-        if not np.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
+        if not qml.math.is_abstract(decomp_mat) and not qml.math.allclose(
+            nphase_mat @ givens_mat.conj(), decomp_mat
+        ):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        unitary[i, i], unitary[j, j] = np.diag(nphase_mat)
-        nleft_givens.append((givens_mat.conj(), (i, j)))
-
-    phases, ordered_rotations = np.diag(unitary), []
-    for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
-        if not np.all(np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])):  # pragma: no cover
-            raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
-        ordered_rotations.append((grot_mat, (i, j)))
-
-    return phases, ordered_rotations
-
-
-def givens_decomposition_jax(unitary):
-    r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.
-
-    This decomposition is based on the construction scheme given in `Optica, 3, 1460 (2016) <https://opg.optica.org/optica/fulltext.cfm?uri=optica-3-12-1460&id=355743>`_\ ,
-    which allows one to write any unitary matrix :math:`U` as:
-
-    .. math::
-
-        U = D \left(\prod_{(m, n) \in G} T_{m, n}(\theta, \phi)\right),
-
-    where :math:`D` is a diagonal phase matrix, :math:`T(\theta, \phi)` is the Givens rotation gates with phase shifts and :math:`G` defines the
-    specific ordered sequence of the Givens rotation gates acting on wires :math:`(m, n)`. The unitary for the :math:`T(\theta, \phi)` gates
-    appearing in the decomposition is of the following form:
-
-    .. math:: T(\theta, \phi) = \begin{bmatrix}
-                                    1 & 0 & 0 & 0 \\
-                                    0 & e^{i \phi} \cos(\theta) & -\sin(\theta) & 0 \\
-                                    0 & e^{i \phi} \sin(\theta) & \cos(\theta) & 0 \\
-                                    0 & 0 & 0 & 1
-                                \end{bmatrix},
-
-    where :math:`\theta \in [0, \pi/2]` is the angle of rotation in the :math:`\{|01\rangle, |10\rangle \}` subspace
-    and :math:`\phi \in [0, 2 \pi]` represents the phase shift at the first wire.
-
-    This function is to be used with JAX interface.
-
-    **Example**
-
-    .. code-block:: python
-
-        unitary = np.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
-                            [-0.21271+0.34938j, -0.38853+0.36497j,  0.61467-0.41317j],
-                            [ 0.41356-0.20765j, -0.00651-0.66689j,  0.32839-0.48293j]])
-
-        phase_mat, ordered_rotations = givens_decomposition(unitary)
-
-    >>> phase_mat
-    tensor([-0.20604358+0.9785369j , -0.82993272+0.55786114j,
-            0.56230612-0.82692833j], requires_grad=True)
-    >>> ordered_rotations
-    [(tensor([[-0.65087861-0.63937521j, -0.40933651-0.j        ],
-            [-0.29201359-0.28685265j,  0.91238348-0.j        ]], requires_grad=True),
-    (0, 1)),
-    (tensor([[ 0.47970366-0.33308926j, -0.8117487 -0.j        ],
-            [ 0.66677093-0.46298215j,  0.5840069 -0.j        ]], requires_grad=True),
-    (1, 2)),
-    (tensor([[ 0.36147547+0.73779454j, -0.57008306-0.j        ],
-            [ 0.2508207 +0.51194108j,  0.82158706-0.j        ]], requires_grad=True),
-    (0, 1))]
-
-    Args:
-        unitary (tensor): unitary matrix on which decomposition will be performed
-
-    Returns:
-        (np.ndarray, list[(np.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
-
-    Raises:
-        ValueError: if the provided matrix is not square.
-
-    .. details::
-        :title: Theory and Pseudocode
-
-        **Givens Rotation**
-
-        Applying the Givens rotation :math:`T(\theta, \phi)` performs the following transformation of the basis states:
-
-        .. math::
-
-            &|00\rangle \mapsto |00\rangle\\
-            &|01\rangle \mapsto e^{i \phi} \cos(\theta) |01\rangle - \sin(\theta) |10\rangle\\
-            &|10\rangle \mapsto e^{i \phi} \sin(\theta) |01\rangle + \cos(\theta) |10\rangle\\
-            &|11\rangle \mapsto |11\rangle.
-
-        **Pseudocode**
-
-        The algorithm that implements the decomposition is the following:
-
-        .. code-block:: python
-
-            def givens_decomposition(U):
-                for i in range(1, N):
-                    if i % 2:
-                        for j in range(0, i):
-                            # Find T^-1(i-j, i-j+1) matrix that nulls element (N-j, i-j) of U
-                            # Update U = U @ T^-1(i-j, i-j+1)
-                    else:
-                        for j in range(1, i):
-                            # Find T(N+j-i-1, N+j-i) matrix that nulls element (N+j-i, j) of U
-                            # Update U = T(N+j-i-1, N+j-i) @ U
-
-    """
-
-    unitary, (M, N) = qml.math.copy(unitary), qml.math.shape(unitary)
-    if M != N:
-        raise ValueError(f"The unitary matrix should be of shape NxN, got {qml.math.shape(unitary)}")
-
-    left_givens, right_givens = [], []
-    for i in range(1, N):
-        if i % 2:
-            for j in range(0, i):
-                indices = [i - j - 1, i - j]
-                grot_mat = _givens_matrix_jax(*unitary[N - j - 1, indices].T, left=True)
-                unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
-                right_givens.append((grot_mat.conj(), indices))
+        if interface == "jax":
+            unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
+            unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
         else:
-            for j in range(1, i + 1):
-                indices = [N + j - i - 2, N + j - i - 1]
-                grot_mat = _givens_matrix_jax(*unitary[indices, j - 1], left=False)
-                unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
-                left_givens.append((grot_mat, indices))
-
-    nleft_givens = []
-    for grot_mat, (i, j) in reversed(left_givens):
-        sphase_mat = qml.math.diag(qml.math.diag(unitary)[qml.math.array([i, j])])
-        decomp_mat = grot_mat.conj().T @ sphase_mat
-        givens_mat = _givens_matrix_jax(*decomp_mat[1, :].T)
-        nphase_mat = decomp_mat @ givens_mat.T
-
-        # check for T_{m,n}^{-1} x D = D x T.
-        if not qml.math.is_abstract(nphase_mat @ givens_mat.conj()) and not qml.math.is_abstract(decomp_mat):
-            if not qml.math.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
-                raise ValueError("Failed to shift phase transposition.")
-
-        unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
-        unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
+            unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
         nleft_givens.append((givens_mat.conj(), (i, j)))
 
     phases, ordered_rotations = qml.math.diag(unitary), []
-    if not qml.math.is_abstract(grot_mat):
-        for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
-            if not qml.math.all(qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])):  # pragma: no cover
-                raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
+    for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
+        if not qml.math.is_abstract(grot_mat) and not qml.math.all(
+            np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])
+        ):  # pragma: no cover
+            raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))
 
     return phases, ordered_rotations
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index f04b7f86bb8..6f75484392b 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -17,7 +17,8 @@
 
 import pennylane as qml
 from pennylane.operation import AnyWires, Operation
-from pennylane.qchem.givens_decomposition import givens_decomposition, givens_decomposition_jax
+from pennylane.qchem.givens_decomposition import givens_decomposition  # , givens_decomposition_jax
+
 
 # pylint: disable-msg=too-many-arguments
 class BasisRotation(Operation):
@@ -116,7 +117,9 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
 
         if check:
             umat = qml.math.copy(unitary_matrix)
-            if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-6):
+            if not qml.math.allclose(
+                umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-6
+            ):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
@@ -159,8 +162,10 @@ def compute_decomposition(
 
         if check:
             umat = qml.math.copy(unitary_matrix)
-            if (not qml.math.is_abstract(unitary_matrix)):
-                if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4):
+            if not qml.math.is_abstract(unitary_matrix):
+                if not qml.math.allclose(
+                    umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4
+                ):
                     raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
@@ -168,10 +173,7 @@ def compute_decomposition(
 
         op_list = []
 
-        phase_list, givens_list = qml.math.where(
-                                    qml.math.get_interface(unitary_matrix) == 'jax',
-                                    givens_decomposition_jax(unitary_matrix),
-                                    givens_decomposition(unitary_matrix))
+        phase_list, givens_list = givens_decomposition(unitary_matrix)
 
         for idx, phase in enumerate(phase_list):
             op_list.append(qml.PhaseShift(qml.math.angle(phase), wires=wires[idx]))
@@ -184,7 +186,7 @@ def compute_decomposition(
                 qml.SingleExcitation(2 * theta, wires=[wires[indices[0]], wires[indices[1]]])
             )
 
-            if not qml.math.isclose(phi, 0.0):
+            if qml.math.is_abstract(phi) or not qml.math.isclose(phi, 0.0):
                 op_list.append(qml.PhaseShift(phi, wires=wires[indices[0]]))
 
         return op_list
diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index 6e9866ad2be..5c62a37598d 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -334,12 +334,12 @@ def test_id(self):
         assert template.id == "a"
 
 
-def circuit_template(unitary_matrix, check = False):
+def circuit_template(unitary_matrix, check=False):
     qml.BasisState(qml.math.array([1, 1, 0]), wires=[0, 1, 2])
     qml.BasisRotation(
         wires=range(3),
         unitary_matrix=unitary_matrix,
-        check = check,
+        check=check,
     )
     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
@@ -402,51 +402,6 @@ def test_autograd(self, tol):
 
         assert np.allclose(grads, np.zeros_like(unitary_matrix, dtype=complex), atol=tol, rtol=0)
 
-    @pytest.mark.jax
-    def test_jax(self, tol):
-        """Test the jax interface."""
-
-        import jax
-        import jax.numpy as jnp
-
-        unitary_matrix = jnp.array(
-            [
-                [0.51378719 + 0.0j, 0.0546265 + 0.79145487j, -0.2051466 + 0.2540723j],
-                [0.62651582 + 0.0j, -0.00828925 - 0.60570321j, -0.36704948 + 0.32528067j],
-                [-0.58608928 + 0.0j, 0.03902657 + 0.04633548j, -0.57220635 + 0.57044649j],
-            ]
-        )
-        weights = jnp.array(
-            [
-                2.2707802713289267,
-                2.9355948424220206,
-                -1.4869222527726533,
-                1.2601662579297865,
-                2.3559705032936717,
-                1.1748572730890159,
-                2.2500537657656356,
-                -0.7251404204443089,
-                2.3577346350335198,
-            ]
-        )
-
-        dev = qml.device("default.qubit", wires=3)
-
-        circuit = qml.QNode(circuit_template, dev)
-        circuit2 = qml.QNode(circuit_decomposed, dev)
-
-        res = circuit(unitary_matrix)
-        res2 = circuit2(weights)
-        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
-
-        grad_fn = jax.grad(circuit)
-        grads = grad_fn(unitary_matrix)
-
-        grad_fn2 = jax.grad(circuit2)
-        grads2 = grad_fn2(weights)
-
-        assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
-
     @pytest.mark.jax
     def test_jax_jit(self, tol):
         """Test the jax interface."""
@@ -461,34 +416,21 @@ def test_jax_jit(self, tol):
                 [-0.58608928 + 0.0j, 0.03902657 + 0.04633548j, -0.57220635 + 0.57044649j],
             ]
         )
-        weights = jnp.array(
-            [
-                2.2707802713289267,
-                2.9355948424220206,
-                -1.4869222527726533,
-                1.2601662579297865,
-                2.3559705032936717,
-                1.1748572730890159,
-                2.2500537657656356,
-                -0.7251404204443089,
-                2.3577346350335198,
-            ]
-        )
 
         dev = qml.device("default.qubit", wires=3)
 
-        circuit = qml.QNode(jax.jit(circuit_template, static_argnames='check'), dev)
-        circuit2 = qml.QNode(circuit_decomposed, dev)
+        circuit = jax.jit(qml.QNode(circuit_template, dev), static_argnames="check")
+        circuit2 = qml.QNode(circuit_template, dev)
 
         res = circuit(unitary_matrix)
-        res2 = circuit2(weights)
+        res2 = circuit2(unitary_matrix)
         assert jnp.allclose(res, res2, atol=tol, rtol=0)
 
         grad_fn = jax.grad(circuit)
         grads = grad_fn(unitary_matrix)
 
         grad_fn2 = jax.grad(circuit2)
-        grads2 = grad_fn2(weights)
+        grads2 = grad_fn2(unitary_matrix)
 
         assert qml.math.allclose(grads[0], grads2[0], atol=tol, rtol=0)
 

From 43686b958b3cfee5f28fa9a3cc07172f9f8b74de Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Tue, 23 Jul 2024 11:19:25 +0200
Subject: [PATCH 24/55] Update basis_rotation.py

---
 pennylane/templates/subroutines/basis_rotation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index f04b7f86bb8..ab3ff632ef2 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -159,7 +159,7 @@ def compute_decomposition(
 
         if check:
             umat = qml.math.copy(unitary_matrix)
-            if (not qml.math.is_abstract(unitary_matrix)):
+            if not qml.math.is_abstract(unitary_matrix):
                 if not qml.math.allclose(umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4):
                     raise ValueError("The provided transformation matrix should be unitary.")
 

From 75edef14d4ec5b8d164b3605bf82d14af01d3399 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Tue, 23 Jul 2024 12:04:00 +0200
Subject: [PATCH 25/55] Update givens_decomposition.py

---
 pennylane/qchem/givens_decomposition.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 8770c469ec2..36866693708 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -14,7 +14,6 @@
 """
 This module contains the functions needed for performing basis transformations defined by a set of fermionic ladder operators.
 """
-import numpy as np
 
 import pennylane as qml
 
@@ -37,7 +36,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         tol (float): determines tolerance limits for :math:`|a|` and :math:`|b|` under which they are considered as zero.
 
     Returns:
-        np.ndarray (or tensor): Givens rotation matrix
+        qml.math.array (or tensor): Givens rotation matrix
 
     """
     abs_a, abs_b, interface = qml.math.abs(a), qml.math.abs(b), qml.math.get_interface(a)
@@ -86,7 +85,7 @@ def givens_decomposition(unitary):
 
     .. code-block:: python
 
-        unitary = np.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
+        unitary = qml.math.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
                             [-0.21271+0.34938j, -0.38853+0.36497j,  0.61467-0.41317j],
                             [ 0.41356-0.20765j, -0.00651-0.66689j,  0.32839-0.48293j]])
 
@@ -110,7 +109,7 @@ def givens_decomposition(unitary):
         unitary (tensor): unitary matrix on which decomposition will be performed
 
     Returns:
-        (np.ndarray, list[(np.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
+        (qml.math.array, list[(qml.math.array, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
 
     Raises:
         ValueError: if the provided matrix is not square.
@@ -207,7 +206,7 @@ def givens_decomposition(U):
     phases, ordered_rotations = qml.math.diag(unitary), []
     for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
         if not qml.math.is_abstract(grot_mat) and not qml.math.all(
-            np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])
+            qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])
         ):  # pragma: no cover
             raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))

From 2c7fffd12edc5f44f76b68db798c04c7ba9926b0 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Tue, 23 Jul 2024 12:41:17 +0200
Subject: [PATCH 26/55] Revert "Update givens_decomposition.py"

This reverts commit 75edef14d4ec5b8d164b3605bf82d14af01d3399.
---
 pennylane/qchem/givens_decomposition.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 36866693708..8770c469ec2 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -14,6 +14,7 @@
 """
 This module contains the functions needed for performing basis transformations defined by a set of fermionic ladder operators.
 """
+import numpy as np
 
 import pennylane as qml
 
@@ -36,7 +37,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         tol (float): determines tolerance limits for :math:`|a|` and :math:`|b|` under which they are considered as zero.
 
     Returns:
-        qml.math.array (or tensor): Givens rotation matrix
+        np.ndarray (or tensor): Givens rotation matrix
 
     """
     abs_a, abs_b, interface = qml.math.abs(a), qml.math.abs(b), qml.math.get_interface(a)
@@ -85,7 +86,7 @@ def givens_decomposition(unitary):
 
     .. code-block:: python
 
-        unitary = qml.math.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
+        unitary = np.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
                             [-0.21271+0.34938j, -0.38853+0.36497j,  0.61467-0.41317j],
                             [ 0.41356-0.20765j, -0.00651-0.66689j,  0.32839-0.48293j]])
 
@@ -109,7 +110,7 @@ def givens_decomposition(unitary):
         unitary (tensor): unitary matrix on which decomposition will be performed
 
     Returns:
-        (qml.math.array, list[(qml.math.array, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
+        (np.ndarray, list[(np.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
 
     Raises:
         ValueError: if the provided matrix is not square.
@@ -206,7 +207,7 @@ def givens_decomposition(U):
     phases, ordered_rotations = qml.math.diag(unitary), []
     for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
         if not qml.math.is_abstract(grot_mat) and not qml.math.all(
-            qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])
+            np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])
         ):  # pragma: no cover
             raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))

From 2b789cc49f6813d04f543289467a28d0378e4b6a Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Tue, 23 Jul 2024 13:15:35 +0200
Subject: [PATCH 27/55] Updating the documentation

---
 pennylane/qchem/givens_decomposition.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 8770c469ec2..25df197c12b 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -14,9 +14,9 @@
 """
 This module contains the functions needed for performing basis transformations defined by a set of fermionic ladder operators.
 """
-import numpy as np
 
 import pennylane as qml
+import numpy as np
 
 
 def _givens_matrix(a, b, left=True, tol=1e-8):
@@ -37,7 +37,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         tol (float): determines tolerance limits for :math:`|a|` and :math:`|b|` under which they are considered as zero.
 
     Returns:
-        np.ndarray (or tensor): Givens rotation matrix
+        qml.math.array (or tensor): Givens rotation matrix
 
     """
     abs_a, abs_b, interface = qml.math.abs(a), qml.math.abs(b), qml.math.get_interface(a)
@@ -86,7 +86,7 @@ def givens_decomposition(unitary):
 
     .. code-block:: python
 
-        unitary = np.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
+        unitary = qml.math.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
                             [-0.21271+0.34938j, -0.38853+0.36497j,  0.61467-0.41317j],
                             [ 0.41356-0.20765j, -0.00651-0.66689j,  0.32839-0.48293j]])
 
@@ -110,7 +110,7 @@ def givens_decomposition(unitary):
         unitary (tensor): unitary matrix on which decomposition will be performed
 
     Returns:
-        (np.ndarray, list[(np.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
+        (qml.math.ndarray, list[(qml.math.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
 
     Raises:
         ValueError: if the provided matrix is not square.

From a8c2d9937f2d4cb85f5949a0368e863317cf6265 Mon Sep 17 00:00:00 2001
From: obliviateandsurrender <utkarshazad98@gmail.com>
Date: Tue, 23 Jul 2024 10:28:25 -0400
Subject: [PATCH 28/55] minor fixes

---
 doc/releases/changelog-dev.md                     |  5 ++++-
 pennylane/qchem/givens_decomposition.py           | 13 ++++++-------
 pennylane/templates/subroutines/basis_rotation.py |  2 +-
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4bd9a7c5931..ee606a94b95 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -53,6 +53,10 @@
 * Molecules and Hamiltonians can now be constructed for all the elements present in the periodic table.
   [(#5821)](https://github.com/PennyLaneAI/pennylane/pull/5821)
 
+* `qml.BasisRotation` and `qml.qchem.givens_decomposition` are now jit compatible.
+  [(#6004)](https://github.com/PennyLaneAI/pennylane/pull/6004)
+
+
 <h4>Community contributions 🥳</h4>
 
 * `DefaultQutritMixed` readout error has been added using parameters `readout_relaxation_probs` and 
@@ -138,7 +142,6 @@
 
 * `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
   [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
-* `qml.BasisRotation` is now qjit compatible. ([#6004](`qml.BasisRotation` is now jit compatible ([#6004](https://github.com/PennyLaneAI/pennylane/issues/6004))))
 
 <h3>Contributors ✍️</h3>
 
diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 25df197c12b..51fb8cdf34b 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -16,7 +16,6 @@
 """
 
 import pennylane as qml
-import numpy as np
 
 
 def _givens_matrix(a, b, left=True, tol=1e-8):
@@ -162,7 +161,7 @@ def givens_decomposition(U):
                     indices = [i - j - 1, i - j]
                     grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
                     unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
-                    right_givens.append((grot_mat.conj(), indices))
+                    right_givens.append((qml.math.conj(grot_mat), indices))
             else:
                 for j in range(1, i + 1):
                     indices = [N + j - i - 2, N + j - i - 1]
@@ -176,7 +175,7 @@ def givens_decomposition(U):
                     indices = [i - j - 1, i - j]
                     grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
                     unitary[:, indices] = unitary[:, indices] @ grot_mat.T
-                    right_givens.append((grot_mat.conj(), indices))
+                    right_givens.append((qml.math.conj(grot_mat), indices))
             else:
                 for j in range(1, i + 1):
                     indices = [N + j - i - 2, N + j - i - 1]
@@ -187,13 +186,13 @@ def givens_decomposition(U):
     nleft_givens = []
     for grot_mat, (i, j) in reversed(left_givens):
         sphase_mat = qml.math.diag(qml.math.diag(unitary)[qml.math.array([i, j])])
-        decomp_mat = grot_mat.conj().T @ sphase_mat
+        decomp_mat = qml.math.conj(grot_mat).T @ sphase_mat
         givens_mat = _givens_matrix(*decomp_mat[1, :].T)
         nphase_mat = decomp_mat @ givens_mat.T
 
         # check for T_{m,n}^{-1} x D = D x T.
         if not qml.math.is_abstract(decomp_mat) and not qml.math.allclose(
-            nphase_mat @ givens_mat.conj(), decomp_mat
+            nphase_mat @ qml.math.conj(givens_mat), decomp_mat
         ):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
@@ -202,12 +201,12 @@ def givens_decomposition(U):
             unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
         else:
             unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
-        nleft_givens.append((givens_mat.conj(), (i, j)))
+        nleft_givens.append((qml.math.conj(givens_mat), (i, j)))
 
     phases, ordered_rotations = qml.math.diag(unitary), []
     for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
         if not qml.math.is_abstract(grot_mat) and not qml.math.all(
-            np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])
+            qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])
         ):  # pragma: no cover
             raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 6f75484392b..314704822ab 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -17,7 +17,7 @@
 
 import pennylane as qml
 from pennylane.operation import AnyWires, Operation
-from pennylane.qchem.givens_decomposition import givens_decomposition  # , givens_decomposition_jax
+from pennylane.qchem.givens_decomposition import givens_decomposition
 
 
 # pylint: disable-msg=too-many-arguments

From 7ef90fb8d8a127015d36762af5a3fd0664eceba9 Mon Sep 17 00:00:00 2001
From: obliviateandsurrender <utkarshazad98@gmail.com>
Date: Tue, 23 Jul 2024 10:34:34 -0400
Subject: [PATCH 29/55] assertion for workflow comparison

---
 tests/templates/test_subroutines/test_basis_rotation.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index 5c62a37598d..aa310f4cdbe 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -424,7 +424,9 @@ def test_jax_jit(self, tol):
 
         res = circuit(unitary_matrix)
         res2 = circuit2(unitary_matrix)
+        res3 = circuit2(qml.math.toarray(unitary_matrix))
         assert jnp.allclose(res, res2, atol=tol, rtol=0)
+        assert qml.math.allclose(res, res3, atol=tol, rtol=0)
 
         grad_fn = jax.grad(circuit)
         grads = grad_fn(unitary_matrix)

From 5cfabafc36ef71a32c936440d30aa5f36c37aa15 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Tue, 23 Jul 2024 16:43:00 +0200
Subject: [PATCH 30/55] Update changelog-dev.md

---
 doc/releases/changelog-dev.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 136d18dfd87..a4c47a83819 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -164,5 +164,6 @@ Austin Huang,
 Christina Lee,
 William Maxwell,
 Vincent Michaud-Rioux,
+Pablo A. Moreno Casares,
 Mudit Pandey,
 Erik Schultheis.

From 9010de10d6fd18ca42893c5ab67b5ab52229ded8 Mon Sep 17 00:00:00 2001
From: Pablo Antonio Moreno Casares <paulmorecasa@gmail.com>
Date: Wed, 24 Jul 2024 17:20:29 +0200
Subject: [PATCH 31/55] Update doc/releases/changelog-dev.md

Co-authored-by: Utkarsh <utkarshazad98@gmail.com>
---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 75200ef419f..49de6ca8265 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -53,7 +53,7 @@
   [(#5821)](https://github.com/PennyLaneAI/pennylane/pull/5821)
 
 * `qml.BasisRotation` and `qml.qchem.givens_decomposition` are now jit compatible.
-  [(#6004)](https://github.com/PennyLaneAI/pennylane/pull/6004)
+  [(#6019)](https://github.com/PennyLaneAI/pennylane/pull/6019)
 
 * The `qubit_observable` function is modified to return an ascending wire order for molecular 
   Hamiltonians.

From 62c4c23e64f88bf8dd3c80e9e8f3ad9259f43f80 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Mon, 29 Jul 2024 13:19:17 +0200
Subject: [PATCH 32/55] Updating docstrings

---
 pennylane/qchem/givens_decomposition.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 51fb8cdf34b..ff07e2075d5 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -36,7 +36,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         tol (float): determines tolerance limits for :math:`|a|` and :math:`|b|` under which they are considered as zero.
 
     Returns:
-        qml.math.array (or tensor): Givens rotation matrix
+        tensor_like: Givens rotation matrix
 
     """
     abs_a, abs_b, interface = qml.math.abs(a), qml.math.abs(b), qml.math.get_interface(a)
@@ -85,7 +85,7 @@ def givens_decomposition(unitary):
 
     .. code-block:: python
 
-        unitary = qml.math.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
+        unitary = np.array([[ 0.73678+0.27511j, -0.5095 +0.10704j, -0.06847+0.32515j],
                             [-0.21271+0.34938j, -0.38853+0.36497j,  0.61467-0.41317j],
                             [ 0.41356-0.20765j, -0.00651-0.66689j,  0.32839-0.48293j]])
 
@@ -109,7 +109,7 @@ def givens_decomposition(unitary):
         unitary (tensor): unitary matrix on which decomposition will be performed
 
     Returns:
-        (qml.math.ndarray, list[(qml.math.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
+        (tensor_like, list[(tensor_like, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
 
     Raises:
         ValueError: if the provided matrix is not square.

From 8b4eef592ed567940039236bf7795cc801f0a399 Mon Sep 17 00:00:00 2001
From: obliviateandsurrender <utkarshazad98@gmail.com>
Date: Tue, 30 Jul 2024 10:28:16 -0400
Subject: [PATCH 33/55] attemp warning fix

---
 pennylane/qchem/givens_decomposition.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index ff07e2075d5..1b02d26fefe 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -40,6 +40,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
 
     """
     abs_a, abs_b, interface = qml.math.abs(a), qml.math.abs(b), qml.math.get_interface(a)
+    aprod = qml.math.nan_to_num(abs_b * abs_a)
     hypot = qml.math.hypot(abs_a, abs_b)
 
     cosine = qml.math.where(abs_a < tol, 1.0, qml.math.where(abs_b < tol, 0.0, abs_b / hypot))
@@ -47,7 +48,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     phase = qml.math.where(
         abs_a < tol,
         1.0,
-        qml.math.where(abs_b < tol, 1.0, 1.0 * b / abs_b * qml.math.conj(a) / abs_a),
+        qml.math.where(abs_b < tol, 1.0, (1.0 * b * qml.math.conj(a)) / (aprod + 1e-15)),
     )
 
     if left:

From 4064fa64d1067f6737329f2b63295c14a11d18cc Mon Sep 17 00:00:00 2001
From: Pablo Antonio Moreno Casares <paulmorecasa@gmail.com>
Date: Wed, 31 Jul 2024 16:56:06 +0200
Subject: [PATCH 34/55] Update pennylane/qchem/givens_decomposition.py

Co-authored-by: Josh Izaac <josh146@gmail.com>
---
 pennylane/qchem/givens_decomposition.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 1b02d26fefe..514f49500b0 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -167,7 +167,7 @@ def givens_decomposition(U):
                 for j in range(1, i + 1):
                     indices = [N + j - i - 2, N + j - i - 1]
                     grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :])
+                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :], indices_are_sorted=True, unique_indices=True)
                     left_givens.append((grot_mat, indices))
     else:
         for i in range(1, N):

From be3df11544f9d7e24a349a8a44cc632a08ace268 Mon Sep 17 00:00:00 2001
From: Pablo Antonio Moreno Casares <paulmorecasa@gmail.com>
Date: Wed, 31 Jul 2024 16:56:19 +0200
Subject: [PATCH 35/55] Update pennylane/qchem/givens_decomposition.py

Co-authored-by: Josh Izaac <josh146@gmail.com>
---
 pennylane/qchem/givens_decomposition.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 514f49500b0..47012fd9a01 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -161,7 +161,7 @@ def givens_decomposition(U):
                 for j in range(0, i):
                     indices = [i - j - 1, i - j]
                     grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T)
+                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T, indices_are_sorted=True, unique_indices=True)
                     right_givens.append((qml.math.conj(grot_mat), indices))
             else:
                 for j in range(1, i + 1):

From 7da2f889aeb1be11404b4d0fe628bbc41bf6ea63 Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Sat, 3 Aug 2024 00:14:21 +0200
Subject: [PATCH 36/55] Update givens_decomposition.py

---
 pennylane/qchem/givens_decomposition.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 47012fd9a01..fb8ccc8ecea 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -161,13 +161,19 @@ def givens_decomposition(U):
                 for j in range(0, i):
                     indices = [i - j - 1, i - j]
                     grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T, indices_are_sorted=True, unique_indices=True)
+                    unitary = unitary.at[:, indices].set(
+                        unitary[:, indices] @ grot_mat.T,
+                        indices_are_sorted=True,
+                        unique_indices=True,
+                    )
                     right_givens.append((qml.math.conj(grot_mat), indices))
             else:
                 for j in range(1, i + 1):
                     indices = [N + j - i - 2, N + j - i - 1]
                     grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :], indices_are_sorted=True, unique_indices=True)
+                    unitary = unitary.at[indices, :].set(
+                        grot_mat @ unitary[indices, :], indices_are_sorted=True, unique_indices=True
+                    )
                     left_givens.append((grot_mat, indices))
     else:
         for i in range(1, N):

From db61a73b15549202a0aecdb9a9317c3776d781ba Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Sat, 3 Aug 2024 00:18:12 +0200
Subject: [PATCH 37/55] Update changelog-dev.md

---
 doc/releases/changelog-dev.md | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6648081e520..c23ae980966 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -302,10 +302,6 @@
 * `qml.lie_closure` works with sums of Paulis.
   [(#6023)](https://github.com/PennyLaneAI/pennylane/pull/6023)
 
-* `qml.BasisRotation` works with qjit.
-  [(#6019)](https://github.com/PennyLaneAI/pennylane/pull/6019)
-
-
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):

From 1c752a1fb033ced651d7ab2ae582203f36738ea5 Mon Sep 17 00:00:00 2001
From: Pablo Antonio Moreno Casares <paulmorecasa@gmail.com>
Date: Sat, 3 Aug 2024 00:20:55 +0200
Subject: [PATCH 38/55] Update doc/releases/changelog-dev.md

Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>
---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index c23ae980966..f363c471dc6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -169,7 +169,7 @@
   Hamiltonians.
   [(#5950)](https://github.com/PennyLaneAI/pennylane/pull/5950)
   
-* `qml.BasisRotation` and `qml.qchem.givens_decomposition` are now jit compatible.
+* `qml.BasisRotation` is now jit compatible.
   [(#6019)](https://github.com/PennyLaneAI/pennylane/pull/6019)
 
 * The `CNOT` operator no longer decomposes to itself. Instead, it raises a `qml.DecompositionUndefinedError`.

From 200481d17f44d36b9179df883981a11bbaa8f6e8 Mon Sep 17 00:00:00 2001
From: obliviateandsurrender <utkarshazad98@gmail.com>
Date: Mon, 5 Aug 2024 12:51:26 -0400
Subject: [PATCH 39/55] code cleanup

---
 pennylane/qchem/givens_decomposition.py       | 73 ++++++++++---------
 .../templates/subroutines/basis_rotation.py   | 19 ++---
 2 files changed, 50 insertions(+), 42 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 47012fd9a01..0889e8453e1 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -57,6 +57,27 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     return qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]], like=interface)
 
 
+def _set_unitary_matrix(unitary_matrix, index, value, like=None):
+    """Set the values in the ``unitary_matrix`` at the specified index.
+
+    Args:
+        unitary_matrix (tensor_like): unitary being modified
+        index (Tuple[int | Ellipsis | List[Int]]): index for the slicing the unitary
+        value (tensor_like): new values for the specified index
+        like (str): interface for the unitary matrix.
+
+    Returns:
+        tensor_like: modified unitary
+    """
+    if like == "jax":
+        return unitary_matrix.at[index[0], index[1]].set(
+            value, indices_are_sorted=True, unique_indices=True
+        )
+
+    unitary_matrix[index[0], index[1]] = value
+    return unitary_matrix
+
+
 # pylint:disable = too-many-branches
 def givens_decomposition(unitary):
     r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.
@@ -155,34 +176,23 @@ def givens_decomposition(U):
         raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")
 
     left_givens, right_givens = [], []
-    if interface == "jax":
-        for i in range(1, N):
-            if i % 2:
-                for j in range(0, i):
-                    indices = [i - j - 1, i - j]
-                    grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T, indices_are_sorted=True, unique_indices=True)
-                    right_givens.append((qml.math.conj(grot_mat), indices))
-            else:
-                for j in range(1, i + 1):
-                    indices = [N + j - i - 2, N + j - i - 1]
-                    grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :], indices_are_sorted=True, unique_indices=True)
-                    left_givens.append((grot_mat, indices))
-    else:
-        for i in range(1, N):
-            if i % 2:
-                for j in range(0, i):
-                    indices = [i - j - 1, i - j]
-                    grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                    unitary[:, indices] = unitary[:, indices] @ grot_mat.T
-                    right_givens.append((qml.math.conj(grot_mat), indices))
-            else:
-                for j in range(1, i + 1):
-                    indices = [N + j - i - 2, N + j - i - 1]
-                    grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                    unitary[indices] = grot_mat @ unitary[indices]
-                    left_givens.append((grot_mat, indices))
+    for i in range(1, N):
+        if i % 2:
+            for j in range(0, i):
+                indices = [i - j - 1, i - j]
+                grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
+                unitary = _set_unitary_matrix(
+                    unitary, (Ellipsis, indices), unitary[:, indices] @ grot_mat.T, like=interface
+                )
+                right_givens.append((qml.math.conj(grot_mat), indices))
+        else:
+            for j in range(1, i + 1):
+                indices = [N + j - i - 2, N + j - i - 1]
+                grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
+                unitary = _set_unitary_matrix(
+                    unitary, (indices, Ellipsis), grot_mat @ unitary[indices, :], like=interface
+                )
+                left_givens.append((grot_mat, indices))
 
     nleft_givens = []
     for grot_mat, (i, j) in reversed(left_givens):
@@ -197,11 +207,8 @@ def givens_decomposition(U):
         ):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        if interface == "jax":
-            unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
-            unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
-        else:
-            unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
+        for diag_idx, diag_val in zip([(i, i), (j, j)], qml.math.diag(nphase_mat)):
+            unitary = _set_unitary_matrix(unitary, diag_idx, diag_val, like=interface)
         nleft_givens.append((qml.math.conj(givens_mat), (i, j)))
 
     phases, ordered_rotations = qml.math.diag(unitary), []
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 314704822ab..fd36d2a55d1 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -116,9 +116,10 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
             )
 
         if check:
-            umat = qml.math.copy(unitary_matrix)
-            if not qml.math.allclose(
-                umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-6
+            if not qml.math.is_abstract(unitary_matrix) and not qml.math.allclose(
+                unitary_matrix @ qml.math.conj(unitary_matrix).T,
+                qml.math.eye(M, dtype=complex),
+                atol=1e-4,
             ):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
@@ -161,12 +162,12 @@ def compute_decomposition(
             )
 
         if check:
-            umat = qml.math.copy(unitary_matrix)
-            if not qml.math.is_abstract(unitary_matrix):
-                if not qml.math.allclose(
-                    umat @ umat.conj().T, qml.math.eye(M, dtype=complex), atol=1e-4
-                ):
-                    raise ValueError("The provided transformation matrix should be unitary.")
+            if not qml.math.is_abstract(unitary_matrix) and not qml.math.allclose(
+                unitary_matrix @ qml.math.conj(unitary_matrix).T,
+                qml.math.eye(M, dtype=complex),
+                atol=1e-4,
+            ):
+                raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
             raise ValueError(f"This template requires at least two wires, got {len(wires)}")

From 24d13c7ccabec0b0ce1311a7511ac85a128748eb Mon Sep 17 00:00:00 2001
From: PabloAMC <paulmorecasa@gmail.com>
Date: Wed, 7 Aug 2024 20:48:27 +0200
Subject: [PATCH 40/55] Update qdrift.py

---
 pennylane/templates/subroutines/qdrift.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/templates/subroutines/qdrift.py b/pennylane/templates/subroutines/qdrift.py
index a65152586e0..9c8198b4753 100644
--- a/pennylane/templates/subroutines/qdrift.py
+++ b/pennylane/templates/subroutines/qdrift.py
@@ -81,7 +81,7 @@ class QDrift(Operation):
     The QDrift subroutine provides a method to approximate the matrix exponential of a Hamiltonian
     expressed as a linear combination of terms which in general do not commute. For the Hamiltonian
     :math:`H = \Sigma_j h_j H_{j}`, the product formula is constructed by random sampling from the
-    terms of the Hamiltonian with the probability :math:`p_j = h_j / \sum_{j} hj` as:
+    terms of the Hamiltonian with the probability :math:`p_j = h_j / \sum_{j} h_{j}` as:
 
     .. math::
 

From 51c63259317ab8930c0ef0e0e620c896c5b59b69 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 10:23:35 -0400
Subject: [PATCH 41/55] reverting changelog

---
 doc/releases/changelog-dev.md | 251 ++++++++++++++++++++++++++--------
 1 file changed, 194 insertions(+), 57 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7c51b80bd11..31e3120d7c6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -3,24 +3,91 @@
 # Release 0.39.0-dev (development release)
 
 <h3>New features since last release</h3>
+
+* Added `process_density_matrix` implementations to 5 `StateMeasurement` subclasses:
+  `ExpVal`, `Var`, `Purity`, `MutualInformation`, and `VnEntropy`.
+  This enables `process_density_matrix` to be an abstract method in `StateMeasurement`,
+  facilitating future support for mixed-state devices and expanded density matrix operations. Also, there is a quick fix for the `np.sqrt` call in the `ProbabilityMP` class to be replaced by `qml.math.sqrt`.
+  [(#6330)](https://github.com/PennyLaneAI/pennylane/pull/6330)
+
+* A new class `MomentumQNGOptimizer` is added. It inherits the basic `QNGOptimizer` class and requires one additional hyperparameter (the momentum coefficient) :math:`0 \leq \rho < 1`, the default value being :math:`\rho=0.9`. For :math:`\rho=0` Momentum-QNG reduces to the basic QNG.
+  [(#6240)](https://github.com/PennyLaneAI/pennylane/pull/6240)
  
+* Function is added for generating the spin Hamiltonian for the
+  [Kitaev](https://arxiv.org/abs/cond-mat/0506438) model on a lattice.
+  [(#6174)](https://github.com/PennyLaneAI/pennylane/pull/6174)
+
+* Function is added for generating the spin Hamiltonians for custom lattices.
+  [(#6226)](https://github.com/PennyLaneAI/pennylane/pull/6226)
+
+* Functions are added for generating spin Hamiltonians for [Emery]
+  (https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.58.2794) and
+  [Haldane](https://journals.aps.org/prl/pdf/10.1103/PhysRevLett.61.2015) models on a lattice.
+  [(#6201)](https://github.com/PennyLaneAI/pennylane/pull/6201/)
+
+* A new `qml.vn_entanglement_entropy` measurement process has been added which measures the
+  Von Neumann entanglement entropy of a quantum state.
+  [(#5911)](https://github.com/PennyLaneAI/pennylane/pull/5911)
+
+* A `has_sparse_matrix` property is added to `Operator` to indicate whether a sparse matrix is defined.
+  [(#6278)](https://github.com/PennyLaneAI/pennylane/pull/6278)
+  [(#6310)](https://github.com/PennyLaneAI/pennylane/pull/6310)
+
 <h3>Improvements 🛠</h3>
 
+* `qml.matrix` now works with empty objects (such as empty tapes, `QNode`s and quantum functions that do
+  not call operations, single operators with empty decompositions).
+  [(#6347)](https://github.com/PennyLaneAI/pennylane/pull/6347)
+  
 * PennyLane is now compatible with NumPy 2.0.
-   [(#6061)](https://github.com/PennyLaneAI/pennylane/pull/6061)
+  [(#6061)](https://github.com/PennyLaneAI/pennylane/pull/6061)
+  [(#6258)](https://github.com/PennyLaneAI/pennylane/pull/6258)
+  [(#6342)](https://github.com/PennyLaneAI/pennylane/pull/6342)
+
+* PennyLane is now compatible with Jax 0.4.28.
+  [(#6255)](https://github.com/PennyLaneAI/pennylane/pull/6255)
 
 * `qml.qchem.excitations` now optionally returns fermionic operators.
   [(#6171)](https://github.com/PennyLaneAI/pennylane/pull/6171)
 
-* The `diagonalize_measurements` transform now uses a more efficient method of diagonalization 
+* The `diagonalize_measurements` transform now uses a more efficient method of diagonalization
   when possible, based on the `pauli_rep` of the relevant observables.
-  [#6113](https://github.com/PennyLaneAI/pennylane/pull/6113/)
+  [(#6113)](https://github.com/PennyLaneAI/pennylane/pull/6113/)
+
+* The `QuantumScript.copy` method now takes `operations`, `measurements`, `shots` and 
+  `trainable_params` as keyword arguments. If any of these are passed when copying a 
+  tape, the specified attributes will replace the copied attributes on the new tape.
+  [(#6285)](https://github.com/PennyLaneAI/pennylane/pull/6285)
+  [(#6363)](https://github.com/PennyLaneAI/pennylane/pull/6363)
 
 * The `Hermitian` operator now has a `compute_sparse_matrix` implementation.
   [(#6225)](https://github.com/PennyLaneAI/pennylane/pull/6225)
 
+* `qml.AmplitudeAmplification` is now compatible with QJIT.
+  [(#6306)](https://github.com/PennyLaneAI/pennylane/pull/6306)
+
+* The quantum arithmetic templates are now QJIT compatible.
+  [(#6307)](https://github.com/PennyLaneAI/pennylane/pull/6307)
+  
+* The `qml.Qubitization` template is now QJIT compatible.
+  [(#6305)](https://github.com/PennyLaneAI/pennylane/pull/6305)
+
+* When an observable is repeated on a tape, `tape.diagonalizing_gates` no longer returns the 
+  diagonalizing gates for each instance of the observable. Instead, the diagonalizing gates of
+  each observable on the tape are included just once.
+  [(#6288)](https://github.com/PennyLaneAI/pennylane/pull/6288)
+
+* The number of diagonalizing gates returned in `qml.specs` now follows the `level` keyword argument 
+  regarding whether the diagonalizing gates are modified by device, instead of always counting 
+  unprocessed diagonalizing gates.
+  [(#6290)](https://github.com/PennyLaneAI/pennylane/pull/6290)
+
 <h4>Capturing and representing hybrid programs</h4>
 
+* `qml.wires.Wires` now accepts JAX arrays as input. Furthermore, a `FutureWarning` is no longer raised in `JAX 0.4.30+`
+  when providing JAX tracers as input to `qml.wires.Wires`.
+  [(#6312)](https://github.com/PennyLaneAI/pennylane/pull/6312)
+
 * Differentiation of hybrid programs via `qml.grad` and `qml.jacobian` can now be captured
   into plxpr. When evaluating a captured `qml.grad` (`qml.jacobian`) instruction, it will
   dispatch to `jax.grad` (`jax.jacobian`), which differs from the Autograd implementation
@@ -36,6 +103,13 @@
   `from pennylane.capture.primitives import *`.
   [(#6129)](https://github.com/PennyLaneAI/pennylane/pull/6129)
 
+* All higher order primitives now use `jax.core.Jaxpr` as metadata instead of sometimes
+  using `jax.core.ClosedJaxpr` and sometimes using `jax.core.Jaxpr`.
+  [(#6319)](https://github.com/PennyLaneAI/pennylane/pull/6319)
+
+* `FermiWord` class now has a method to apply anti-commutator relations.
+   [(#6196)](https://github.com/PennyLaneAI/pennylane/pull/6196)
+
 * `FermiWord` and `FermiSentence` classes now have methods to compute adjoints.
   [(#6166)](https://github.com/PennyLaneAI/pennylane/pull/6166)
 
@@ -50,53 +124,34 @@
   unique representation of the object.
   [(#6167)](https://github.com/PennyLaneAI/pennylane/pull/6167)
 
-* If the conditional does not include a mid-circuit measurement, then `qml.cond`
-  will automatically evaluate conditionals using standard Python control flow.
-  [(#6016)](https://github.com/PennyLaneAI/pennylane/pull/6016)
-
-  This allows `qml.cond` to be used to represent a wider range of conditionals:
-
-  ```python
-  dev = qml.device("default.qubit", wires=1)
-
-  @qml.qnode(dev)
-  def circuit(x):
-      c = qml.cond(x > 2.7, qml.RX, qml.RZ)
-      c(x, wires=0)
-      return qml.probs(wires=0)
-  ```
-
-  ```pycon
-  >>> print(qml.draw(circuit)(3.8))
-  0: ──RX(3.80)─┤  Probs
-  >>> print(qml.draw(circuit)(0.54))
-  0: ──RZ(0.54)─┤  Probs
-  ```
+* Predefined lattice shapes such as `lieb`, `cubic`, `bcc`, `fcc`, and `diamond`
+  can now be generated.
+  [(6237)](https://github.com/PennyLaneAI/pennylane/pull/6237)
 
-* The `qubit_observable` function is modified to return an ascending wire order for molecular 
-  Hamiltonians.
-  [(#5950)](https://github.com/PennyLaneAI/pennylane/pull/5950)
-  
-* `qml.BasisRotation` is now jit compatible.
-  [(#6019)](https://github.com/PennyLaneAI/pennylane/pull/6019)
-
-* The `CNOT` operator no longer decomposes to itself. Instead, it raises a `qml.DecompositionUndefinedError`.
-  [(#6039)](https://github.com/PennyLaneAI/pennylane/pull/6039)
-
-<h4>Community contributions 🥳</h4>
-
-* `DefaultQutritMixed` readout error has been added using parameters `readout_relaxation_probs` and 
-  `readout_misclassification_probs` on the `default.qutrit.mixed` device. These parameters add a `~.QutritAmplitudeDamping`  and a `~.TritFlip` channel, respectively,
-  after measurement diagonalization. The amplitude damping error represents the potential for
-  relaxation to occur during longer measurements. The trit flip error represents misclassification during readout.
-  [(#5842)](https://github.com/PennyLaneAI/pennylane/pull/5842)
+* A `ReferenceQubit` is introduced for testing purposes and as a reference for future plugin development.
+  [(#6181)](https://github.com/PennyLaneAI/pennylane/pull/6181)
 
 * The `to_mat` methods for `FermiWord` and `FermiSentence` now optionally return
   a sparse matrix.
   [(#6173)](https://github.com/PennyLaneAI/pennylane/pull/6173)
 
+* The `make_plxpr` function is added, to take a function and create a `Callable` that,
+  when called, will return a PLxPR representation of the input function.
+  [(#6326)](https://github.com/PennyLaneAI/pennylane/pull/6326)
+
 <h3>Breaking changes 💔</h3>
 
+* `AllWires` validation in `QNode.construct` has been removed. 
+  [(#6373)](https://github.com/PennyLaneAI/pennylane/pull/6373)
+
+* The `simplify` argument in `qml.Hamiltonian` and `qml.ops.LinearCombination` has been removed.
+  Instead, `qml.simplify()` can be called on the constructed operator.
+  [(#6279)](https://github.com/PennyLaneAI/pennylane/pull/6279)
+
+* The functions `qml.qinfo.classical_fisher` and `qml.qinfo.quantum_fisher` have been removed and migrated to the `qml.gradients`
+  module. Therefore, `qml.gradients.classical_fisher` and `qml.gradients.quantum_fisher` should be used instead.
+  [(#5911)](https://github.com/PennyLaneAI/pennylane/pull/5911)
+
 * Remove support for Python 3.9.
   [(#6223)](https://github.com/PennyLaneAI/pennylane/pull/6223)
 
@@ -121,12 +176,50 @@
   Please use `qml.transforms.split_non_commuting` instead.
   [(#6204)](https://github.com/PennyLaneAI/pennylane/pull/6204)
 
+* The `decomp_depth` keyword argument to `qml.device` is removed.
+  [(#6234)](https://github.com/PennyLaneAI/pennylane/pull/6234)
+
 * `Operator.expand` is now removed. Use `qml.tape.QuantumScript(op.deocomposition())` instead.
   [(#6227)](https://github.com/PennyLaneAI/pennylane/pull/6227)
 
-
 <h3>Deprecations 👋</h3>
 
+* Legacy operator arithmetic has been deprecated. This includes `qml.ops.Hamiltonian`, `qml.operation.Tensor`,
+  `qml.operation.enable_new_opmath`, `qml.operation.disable_new_opmath`, and `qml.operation.convert_to_legacy_H`.
+  Note that when new operator arithmetic is enabled, ``qml.Hamiltonian`` will continue to dispatch to
+  `qml.ops.LinearCombination`; this behaviour is not deprecated. For more information, check out the
+  [updated operator troubleshooting page](https://docs.pennylane.ai/en/stable/news/new_opmath.html).
+  [(#6287)](https://github.com/PennyLaneAI/pennylane/pull/6287)
+  [(#6365)](https://github.com/PennyLaneAI/pennylane/pull/6365)
+
+* `qml.pauli.PauliSentence.hamiltonian` and `qml.pauli.PauliWord.hamiltonian` are deprecated. Instead, please use
+  `qml.pauli.PauliSentence.operation` and `qml.pauli.PauliWord.operation` respectively.
+  [(#6287)](https://github.com/PennyLaneAI/pennylane/pull/6287)
+
+* `qml.pauli.simplify()` is deprecated. Instead, please use `qml.simplify(op)` or `op.simplify()`.
+  [(#6287)](https://github.com/PennyLaneAI/pennylane/pull/6287)
+
+* The `qml.BasisStatePreparation` template is deprecated.
+  Instead, use `qml.BasisState`.
+  [(#6021)](https://github.com/PennyLaneAI/pennylane/pull/6021)
+
+* The `'ancilla'` argument for `qml.iterative_qpe` has been deprecated. Instead, use the `'aux_wire'` argument.
+  [(#6277)](https://github.com/PennyLaneAI/pennylane/pull/6277)
+
+* `qml.shadows.shadow_expval` has been deprecated. Instead, use the `qml.shadow_expval` measurement
+  process.
+  [(#6277)](https://github.com/PennyLaneAI/pennylane/pull/6277)
+
+* `qml.broadcast` has been deprecated. Please use `for` loops instead.
+  [(#6277)](https://github.com/PennyLaneAI/pennylane/pull/6277)
+
+* The `qml.QubitStateVector` template is deprecated. Instead, use `qml.StatePrep`.
+  [(#6172)](https://github.com/PennyLaneAI/pennylane/pull/6172)
+
+* The `qml.qinfo` module has been deprecated. Please see the respective functions in the `qml.math` and
+  `qml.measurements` modules instead.
+  [(#5911)](https://github.com/PennyLaneAI/pennylane/pull/5911)
+
 * `Device`, `QubitDevice`, and `QutritDevice` will no longer be accessible via top-level import in v0.40.
   They will still be accessible as `qml.devices.LegacyDevice`, `qml.devices.QubitDevice`, and `qml.devices.QutritDevice`
   respectively.
@@ -137,27 +230,57 @@
 
 <h3>Documentation 📝</h3>
 
+* Update `qml.Qubitization` documentation based on new decomposition.
+  [(#6276)](https://github.com/PennyLaneAI/pennylane/pull/6276)
+
+* Fixed examples in the documentation of a few optimizers.
+  [(#6303)](https://github.com/PennyLaneAI/pennylane/pull/6303)
+  [(#6315)](https://github.com/PennyLaneAI/pennylane/pull/6315)
+
+* Corrected examples in the documentation of `qml.jacobian`.
+  [(#6283)](https://github.com/PennyLaneAI/pennylane/pull/6283)
+  [(#6315)](https://github.com/PennyLaneAI/pennylane/pull/6315)
+
+* Fixed spelling in a number of places across the documentation.
+  [(#6280)](https://github.com/PennyLaneAI/pennylane/pull/6280)
+
+* Add `work_wires` parameter to `qml.MultiControlledX` docstring signature.
+  [(#6271)](https://github.com/PennyLaneAI/pennylane/pull/6271)
+
+* Removed ambiguity in error raised by the `PauliRot` class.
+  [(#6298)](https://github.com/PennyLaneAI/pennylane/pull/6298)
+
 <h3>Bug fixes 🐛</h3>
 
+* `adjoint_metric_tensor` now works with circuits containing state preparation operations.
+  [(#6358)](https://github.com/PennyLaneAI/pennylane/pull/6358)
+
+* `quantum_fisher` now respects the classical Jacobian of QNodes.
+  [(#6350)](https://github.com/PennyLaneAI/pennylane/pull/6350)
+
+* `qml.map_wires` can now be applied to a batch of tapes.
+  [(#6295)](https://github.com/PennyLaneAI/pennylane/pull/6295)
+
+* Fix float-to-complex casting in various places across PennyLane.
+ [(#6260)](https://github.com/PennyLaneAI/pennylane/pull/6260)
+ [(#6268)](https://github.com/PennyLaneAI/pennylane/pull/6268)
+
 * Fix a bug where zero-valued JVPs were calculated wrongly in the presence of shot vectors.
   [(#6219)](https://github.com/PennyLaneAI/pennylane/pull/6219)
 
-* Fix `qml.PrepSelPrep` template to work with `torch`:
+* Fix `qml.PrepSelPrep` template to work with `torch`.
   [(#6191)](https://github.com/PennyLaneAI/pennylane/pull/6191)
 
 * Now `qml.equal` compares correctly `qml.PrepSelPrep` operators.
   [(#6182)](https://github.com/PennyLaneAI/pennylane/pull/6182)
 
-* The ``qml.QSVT`` template now orders the ``projector`` wires first and the ``UA`` wires second, which is the expected order of the decomposition.
+* The `qml.QSVT` template now orders the `projector` wires first and the `UA` wires second, which is the expected order of the decomposition.
   [(#6212)](https://github.com/PennyLaneAI/pennylane/pull/6212)
-  
-* The ``qml.Qubitization`` template now orders the ``control`` wires first and the ``hamiltonian`` wires second, which is the expected according to other templates.
-  [(#6229)](https://github.com/PennyLaneAI/pennylane/pull/6229)
 
-* `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
-  [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
+* The `qml.Qubitization` template now orders the `control` wires first and the `hamiltonian` wires second, which is the expected according to other templates.
+  [(#6229)](https://github.com/PennyLaneAI/pennylane/pull/6229)
 
-* The ``qml.FABLE`` template now returns the correct value when JIT is enabled.
+* The `qml.FABLE` template now returns the correct value when JIT is enabled.
   [(#6263)](https://github.com/PennyLaneAI/pennylane/pull/6263)
 
 * Fixes a bug where a circuit using the `autograd` interface sometimes returns nested values that are not of the `autograd` interface.
@@ -166,23 +289,37 @@
 * Fixes a bug where a simple circuit with no parameters or only builtin/numpy arrays as parameters returns autograd tensors.
   [(#6225)](https://github.com/PennyLaneAI/pennylane/pull/6225)
 
+* `qml.pauli.PauliVSpace` now uses a more stable SVD-based linear independence check to avoid running into `LinAlgError: Singular matrix`. This stabilizes the usage of `qml.lie_closure`. It also introduces normalization of the basis vector's internal representation `_M` to avoid exploding coefficients.
+  [(#6232)](https://github.com/PennyLaneAI/pennylane/pull/6232)
+
+* Fixes a bug where `csc_dot_product` is used during measurement for `Sum`/`Hamiltonian` that contains observables that does not define a sparse matrix.
+  [(#6278)](https://github.com/PennyLaneAI/pennylane/pull/6278)
+  [(#6310)](https://github.com/PennyLaneAI/pennylane/pull/6310)
+
+* Fixes a test after updating to the nightly version of Catalyst.
+  [(#6362)](https://github.com/PennyLaneAI/pennylane/pull/6362)
+
+* Fixes a bug where `CommutingEvolution` with a trainable `Hamiltonian` cannot be differentiated using parameter shift.
+  [(#6372)](https://github.com/PennyLaneAI/pennylane/pull/6372)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
 
 Guillermo Alonso,
 Utkarsh Azad,
+Oleksandr Borysenko,
 Astral Cai,
+Isaac De Vlugt,
+Diksha Dhawan,
 Lillian M. A. Frederiksen,
 Pietropaolo Frisoni,
 Emiliano Godinez,
+Austin Huang,
+Korbinian Kottmann,
 Christina Lee,
 William Maxwell,
+Erick Ochoa Lopez,
 Lee J. O'Riordan,
-Vincent Michaud-Rioux,
-Anurav Modak,
-Pablo A. Moreno Casares,
 Mudit Pandey,
-Erik Schultheis,
-Nate Stemen,
-David Wierichs.
\ No newline at end of file
+David Wierichs,

From b3b0bbaa6b14c80f059ce93f8e276c6123cd8fe8 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 10:26:14 -0400
Subject: [PATCH 42/55] reverting qdrift changes

---
 pennylane/templates/subroutines/qdrift.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/templates/subroutines/qdrift.py b/pennylane/templates/subroutines/qdrift.py
index 9c8198b4753..a65152586e0 100644
--- a/pennylane/templates/subroutines/qdrift.py
+++ b/pennylane/templates/subroutines/qdrift.py
@@ -81,7 +81,7 @@ class QDrift(Operation):
     The QDrift subroutine provides a method to approximate the matrix exponential of a Hamiltonian
     expressed as a linear combination of terms which in general do not commute. For the Hamiltonian
     :math:`H = \Sigma_j h_j H_{j}`, the product formula is constructed by random sampling from the
-    terms of the Hamiltonian with the probability :math:`p_j = h_j / \sum_{j} h_{j}` as:
+    terms of the Hamiltonian with the probability :math:`p_j = h_j / \sum_{j} hj` as:
 
     .. math::
 

From 99ee770c11ff81e743596d5e0dc13087f21538d3 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 10:44:00 -0400
Subject: [PATCH 43/55] cleanup givens changes

---
 pennylane/qchem/givens_decomposition.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 0889e8453e1..241c21a419f 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -43,13 +43,14 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     aprod = qml.math.nan_to_num(abs_b * abs_a)
     hypot = qml.math.hypot(abs_a, abs_b)
 
-    cosine = qml.math.where(abs_a < tol, 1.0, qml.math.where(abs_b < tol, 0.0, abs_b / hypot))
-    sine = qml.math.where(abs_a < tol, 0.0, qml.math.where(abs_b < tol, 1.0, abs_a / hypot))
-    phase = qml.math.where(
-        abs_a < tol,
-        1.0,
-        qml.math.where(abs_b < tol, 1.0, (1.0 * b * qml.math.conj(a)) / (aprod + 1e-15)),
-    )
+    cosine = qml.math.where(abs_b < tol, 0.0, abs_b / hypot)
+    cosine = qml.math.where(abs_a < tol, 1.0, cosine)
+
+    sine = qml.math.where(abs_b < tol, 1.0, abs_a / hypot)
+    sine = qml.math.where(abs_a < tol, 0.0, sine)
+
+    phase = qml.math.where(abs_b < tol, 1.0, (1.0 * b * qml.math.conj(a)) / (aprod+1e-15))
+    phase = qml.math.where(abs_a < tol, 0.0, phase)
 
     if left:
         return qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]], like=interface)

From 443017e725d17fd7b4709c7243a4ca48367ac19f Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 10:45:59 -0400
Subject: [PATCH 44/55] update changelog

---
 doc/releases/changelog-dev.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1804494a8ae..07e62368fc5 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -85,6 +85,9 @@
   unprocessed diagonalizing gates.
   [(#6290)](https://github.com/PennyLaneAI/pennylane/pull/6290)
 
+* `qml.BasisRotation` template is now JIT compatible.
+  [(#6019)](https://github.com/PennyLaneAI/pennylane/pull/6019)
+
 <h4>Capturing and representing hybrid programs</h4>
 
 * `qml.wires.Wires` now accepts JAX arrays as input. Furthermore, a `FutureWarning` is no longer raised in `JAX 0.4.30+`

From 8192fa05f4ba30f92ec9a9929f324a8ed97739a8 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 11:05:48 -0400
Subject: [PATCH 45/55] minor changes to test

---
 tests/templates/test_subroutines/test_basis_rotation.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index aa310f4cdbe..05556641114 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -425,7 +425,8 @@ def test_jax_jit(self, tol):
         res = circuit(unitary_matrix)
         res2 = circuit2(unitary_matrix)
         res3 = circuit2(qml.math.toarray(unitary_matrix))
-        assert jnp.allclose(res, res2, atol=tol, rtol=0)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
         assert qml.math.allclose(res, res3, atol=tol, rtol=0)
 
         grad_fn = jax.grad(circuit)
@@ -434,7 +435,7 @@ def test_jax_jit(self, tol):
         grad_fn2 = jax.grad(circuit2)
         grads2 = grad_fn2(unitary_matrix)
 
-        assert qml.math.allclose(grads[0], grads2[0], atol=tol, rtol=0)
+        assert qml.math.allclose(grads, grads2, atol=tol, rtol=0)
 
     @pytest.mark.tf
     def test_tf(self, tol):

From e4cc0d2c0cce8ac2e3f349bf3e4170021c8ec18f Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 11:06:26 -0400
Subject: [PATCH 46/55] black

---
 pennylane/qchem/givens_decomposition.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 241c21a419f..54863c0d1f8 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -49,7 +49,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     sine = qml.math.where(abs_b < tol, 1.0, abs_a / hypot)
     sine = qml.math.where(abs_a < tol, 0.0, sine)
 
-    phase = qml.math.where(abs_b < tol, 1.0, (1.0 * b * qml.math.conj(a)) / (aprod+1e-15))
+    phase = qml.math.where(abs_b < tol, 1.0, (1.0 * b * qml.math.conj(a)) / (aprod + 1e-15))
     phase = qml.math.where(abs_a < tol, 0.0, phase)
 
     if left:

From 18dd017a3727f0ffde3a2da9d3678bf047c31c15 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 11:42:29 -0400
Subject: [PATCH 47/55] fixing minor bug

---
 pennylane/qchem/givens_decomposition.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 54863c0d1f8..6abfd2d2e84 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -50,7 +50,7 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     sine = qml.math.where(abs_a < tol, 0.0, sine)
 
     phase = qml.math.where(abs_b < tol, 1.0, (1.0 * b * qml.math.conj(a)) / (aprod + 1e-15))
-    phase = qml.math.where(abs_a < tol, 0.0, phase)
+    phase = qml.math.where(abs_a < tol, 1.0, phase)
 
     if left:
         return qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]], like=interface)

From a11d2db74d7b59d178b834fdc7ad4aeeb5d6339b Mon Sep 17 00:00:00 2001
From: Utkarsh <utkarshazad98@gmail.com>
Date: Fri, 11 Oct 2024 12:57:50 -0400
Subject: [PATCH 48/55] Update
 pennylane/templates/subroutines/basis_rotation.py

---
 pennylane/templates/subroutines/basis_rotation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index fd36d2a55d1..a800cecf4de 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -112,7 +112,7 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
 
         if M != N:
             raise ValueError(
-                f"The unitary matrix should be of shape NxN, got {qml.math.shape(unitary_matrix)}"
+                f"The unitary matrix should be of shape NxN, got {(M, N)}"
             )
 
         if check:

From 2dcaf6c56731c528434a7ee4a788c9bca9a5d0ca Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Fri, 11 Oct 2024 13:09:22 -0400
Subject: [PATCH 49/55] black

---
 pennylane/templates/subroutines/basis_rotation.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index a800cecf4de..832a0684739 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -111,9 +111,7 @@ def __init__(self, wires, unitary_matrix, check=False, id=None):
         M, N = qml.math.shape(unitary_matrix)
 
         if M != N:
-            raise ValueError(
-                f"The unitary matrix should be of shape NxN, got {(M, N)}"
-            )
+            raise ValueError(f"The unitary matrix should be of shape NxN, got {(M, N)}")
 
         if check:
             if not qml.math.is_abstract(unitary_matrix) and not qml.math.allclose(

From 80e6654b8d187fba8fac0f57b0548d3873313102 Mon Sep 17 00:00:00 2001
From: Utkarsh <utkarshazad98@gmail.com>
Date: Tue, 15 Oct 2024 20:41:18 -0400
Subject: [PATCH 50/55] Apply suggestions from code review

Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>
---
 pennylane/qchem/givens_decomposition.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 6abfd2d2e84..0847920f1aa 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -63,9 +63,9 @@ def _set_unitary_matrix(unitary_matrix, index, value, like=None):
 
     Args:
         unitary_matrix (tensor_like): unitary being modified
-        index (Tuple[int | Ellipsis | List[Int]]): index for the slicing the unitary
+        index (Tuple[Int | Ellipsis | List[Int]]): index for slicing the unitary
         value (tensor_like): new values for the specified index
-        like (str): interface for the unitary matrix.
+        like (str): interface for the unitary matrix
 
     Returns:
         tensor_like: modified unitary
@@ -132,7 +132,7 @@ def givens_decomposition(unitary):
         unitary (tensor): unitary matrix on which decomposition will be performed
 
     Returns:
-        (tensor_like, list[(tensor_like, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
+        (tensor_like, list[(tensor_like, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices
 
     Raises:
         ValueError: if the provided matrix is not square.

From 2b01f241a2ae75fdf50cbbdbe4967f968129cee6 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Wed, 16 Oct 2024 12:43:13 -0400
Subject: [PATCH 51/55] addressing soran's comments

---
 pennylane/qchem/givens_decomposition.py | 24 +++++++++++++++++++++++-
 tests/qchem/test_givens_rotations.py    | 12 ++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 0847920f1aa..b09d18ff91e 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -38,8 +38,21 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
     Returns:
         tensor_like: Givens rotation matrix
 
+    Raises:
+        TypeError: if a and b have different interfaces
+
     """
-    abs_a, abs_b, interface = qml.math.abs(a), qml.math.abs(b), qml.math.get_interface(a)
+
+    abs_a, abs_b = qml.math.abs(a), qml.math.abs(b)
+    interface_a, interface_b = qml.math.get_interface(a), qml.math.get_interface(b)
+
+    if interface_a != interface_b:
+        raise TypeError(
+            f"The interfaces of 'a' and 'b' do not match. Got {interface_a} and {interface_b}."
+        )
+
+    interface = interface_a
+
     aprod = qml.math.nan_to_num(abs_b * abs_a)
     hypot = qml.math.hypot(abs_a, abs_b)
 
@@ -69,7 +82,16 @@ def _set_unitary_matrix(unitary_matrix, index, value, like=None):
 
     Returns:
         tensor_like: modified unitary
+
+    Examples:
+        A = np.eye(5)
+        A = _set_unitary_matrix(A, (0, 0), 5)
+        A = _set_unitary_matrix(A, (1, Ellipsis), np.array([1, 2, 3, 4, 5]))
+        A = _set_unitary_matrix(A, (1, [1, 2]), np.array([3, 4]))
     """
+    if like is None:
+        like = qml.math.get_interface(unitary_matrix)
+
     if like == "jax":
         return unitary_matrix.at[index[0], index[1]].set(
             value, indices_are_sorted=True, unique_indices=True
diff --git a/tests/qchem/test_givens_rotations.py b/tests/qchem/test_givens_rotations.py
index ebb254cc72f..1bef9242545 100644
--- a/tests/qchem/test_givens_rotations.py
+++ b/tests/qchem/test_givens_rotations.py
@@ -126,3 +126,15 @@ def test_givens_decomposition_exceptions(unitary_matrix, msg_match):
 
     with pytest.raises(ValueError, match=msg_match):
         givens_decomposition(unitary_matrix)
+
+
+@pytest.mark.jax
+def test_givens_matrix_exceptions():
+    """Test that _givens_matrix throws an exception if the parameters have different interface."""
+    import jax.numpy as jnp
+
+    a = np.array(1.2)
+    b = jnp.array(2.3)
+
+    with pytest.raises(TypeError, match="The interfaces of 'a' and 'b' do not match."):
+        _givens_matrix(a, b)

From 63f2ea7ea8060f79bc21d63b2c7d76393ece3399 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Wed, 16 Oct 2024 14:19:50 -0400
Subject: [PATCH 52/55] test for _set_untiary_matrix

---
 tests/qchem/test_givens_rotations.py | 34 +++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/tests/qchem/test_givens_rotations.py b/tests/qchem/test_givens_rotations.py
index 1bef9242545..14f5b80cb9e 100644
--- a/tests/qchem/test_givens_rotations.py
+++ b/tests/qchem/test_givens_rotations.py
@@ -18,8 +18,15 @@
 import pytest
 from scipy.stats import unitary_group
 
+import pennylane as qml
 from pennylane import numpy as np
-from pennylane.qchem.givens_decomposition import _givens_matrix, givens_decomposition
+from pennylane.qchem.givens_decomposition import (
+    _givens_matrix,
+    _set_unitary_matrix,
+    givens_decomposition,
+)
+
+jnp = pytest.importorskip("jax.numpy")
 
 
 @pytest.mark.parametrize("left", [True, False])
@@ -138,3 +145,28 @@ def test_givens_matrix_exceptions():
 
     with pytest.raises(TypeError, match="The interfaces of 'a' and 'b' do not match."):
         _givens_matrix(a, b)
+
+
+@pytest.mark.parametrize(
+    ("unitary_matrix", "index", "value", "like", "expected_matrix"),
+    [
+        (np.array([[1, 0], [0, 1]]), (0, 0), 5, None, np.array([[5, 0], [0, 1]])),
+        (np.array([[1, 0], [0, 1]]), (0, 0), 5, "numpy", np.array([[5, 0], [0, 1]])),
+        (np.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], None, np.array([[1, 2], [0, 1]])),
+        (np.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], "numpy", np.array([[1, 2], [0, 1]])),
+        (np.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], None, np.array([[1, 0], [1, 2]])),
+        (np.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], "numpy", np.array([[1, 0], [1, 2]])),
+        (jnp.array([[1, 0], [0, 1]]), (0, 0), 5, None, jnp.array([[5, 0], [0, 1]])),
+        (jnp.array([[1, 0], [0, 1]]), (0, 0), 5, "jax", jnp.array([[5, 0], [0, 1]])),
+        (jnp.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], None, jnp.array([[1, 2], [0, 1]])),
+        (jnp.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], "jax", jnp.array([[1, 2], [0, 1]])),
+        (jnp.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], None, jnp.array([[1, 0], [1, 2]])),
+        (jnp.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], "jax", jnp.array([[1, 0], [1, 2]])),
+    ],
+)
+@pytest.mark.jax
+def test_set_unitary_matrix(unitary_matrix, index, value, like, expected_matrix):
+    """Test the _set_unitary function on different interfaces."""
+
+    new_unitary_matrix = _set_unitary_matrix(unitary_matrix, index, value, like)
+    assert qml.math.allclose(new_unitary_matrix, expected_matrix)

From c45986259d7955db30246c25f16ec0dfdd9f99c8 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Wed, 16 Oct 2024 14:33:03 -0400
Subject: [PATCH 53/55] linting

---
 tests/qchem/test_givens_rotations.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/qchem/test_givens_rotations.py b/tests/qchem/test_givens_rotations.py
index 14f5b80cb9e..1718c603c91 100644
--- a/tests/qchem/test_givens_rotations.py
+++ b/tests/qchem/test_givens_rotations.py
@@ -138,7 +138,6 @@ def test_givens_decomposition_exceptions(unitary_matrix, msg_match):
 @pytest.mark.jax
 def test_givens_matrix_exceptions():
     """Test that _givens_matrix throws an exception if the parameters have different interface."""
-    import jax.numpy as jnp
 
     a = np.array(1.2)
     b = jnp.array(2.3)

From e10576bde00b72a98a8cad702e71fdd35c32e6e5 Mon Sep 17 00:00:00 2001
From: Will <wmaxwell90@gmail.com>
Date: Wed, 16 Oct 2024 15:26:30 -0400
Subject: [PATCH 54/55] removing jax import from test

---
 tests/qchem/test_givens_rotations.py | 57 ++++++++++++++++++++--------
 1 file changed, 41 insertions(+), 16 deletions(-)

diff --git a/tests/qchem/test_givens_rotations.py b/tests/qchem/test_givens_rotations.py
index 1718c603c91..b28b78c9143 100644
--- a/tests/qchem/test_givens_rotations.py
+++ b/tests/qchem/test_givens_rotations.py
@@ -26,8 +26,6 @@
     givens_decomposition,
 )
 
-jnp = pytest.importorskip("jax.numpy")
-
 
 @pytest.mark.parametrize("left", [True, False])
 @pytest.mark.parametrize(
@@ -138,6 +136,7 @@ def test_givens_decomposition_exceptions(unitary_matrix, msg_match):
 @pytest.mark.jax
 def test_givens_matrix_exceptions():
     """Test that _givens_matrix throws an exception if the parameters have different interface."""
+    import jax.numpy as jnp
 
     a = np.array(1.2)
     b = jnp.array(2.3)
@@ -147,25 +146,51 @@ def test_givens_matrix_exceptions():
 
 
 @pytest.mark.parametrize(
-    ("unitary_matrix", "index", "value", "like", "expected_matrix"),
+    ("jax", "unitary_matrix", "index", "value", "like", "expected_matrix"),
     [
-        (np.array([[1, 0], [0, 1]]), (0, 0), 5, None, np.array([[5, 0], [0, 1]])),
-        (np.array([[1, 0], [0, 1]]), (0, 0), 5, "numpy", np.array([[5, 0], [0, 1]])),
-        (np.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], None, np.array([[1, 2], [0, 1]])),
-        (np.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], "numpy", np.array([[1, 2], [0, 1]])),
-        (np.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], None, np.array([[1, 0], [1, 2]])),
-        (np.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], "numpy", np.array([[1, 0], [1, 2]])),
-        (jnp.array([[1, 0], [0, 1]]), (0, 0), 5, None, jnp.array([[5, 0], [0, 1]])),
-        (jnp.array([[1, 0], [0, 1]]), (0, 0), 5, "jax", jnp.array([[5, 0], [0, 1]])),
-        (jnp.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], None, jnp.array([[1, 2], [0, 1]])),
-        (jnp.array([[1, 0], [0, 1]]), (0, Ellipsis), [1, 2], "jax", jnp.array([[1, 2], [0, 1]])),
-        (jnp.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], None, jnp.array([[1, 0], [1, 2]])),
-        (jnp.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], "jax", jnp.array([[1, 0], [1, 2]])),
+        (False, np.array([[1, 0], [0, 1]]), (0, 0), 5, None, np.array([[5, 0], [0, 1]])),
+        (False, np.array([[1, 0], [0, 1]]), (0, 0), 5, "numpy", np.array([[5, 0], [0, 1]])),
+        (
+            False,
+            np.array([[1, 0], [0, 1]]),
+            (0, Ellipsis),
+            [1, 2],
+            None,
+            np.array([[1, 2], [0, 1]]),
+        ),
+        (
+            False,
+            np.array([[1, 0], [0, 1]]),
+            (0, Ellipsis),
+            [1, 2],
+            "numpy",
+            np.array([[1, 2], [0, 1]]),
+        ),
+        (False, np.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], None, np.array([[1, 0], [1, 2]])),
+        (
+            False,
+            np.array([[1, 0], [0, 1]]),
+            (1, [0, 1]),
+            [1, 2],
+            "numpy",
+            np.array([[1, 0], [1, 2]]),
+        ),
+        (True, [[1, 0], [0, 1]], (0, 0), 5, None, [[5, 0], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (0, 0), 5, "jax", [[5, 0], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (0, Ellipsis), [1, 2], None, [[1, 2], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (0, Ellipsis), [1, 2], "jax", [[1, 2], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (1, [0, 1]), [1, 2], None, [[1, 0], [1, 2]]),
+        (True, [[1, 0], [0, 1]], (1, [0, 1]), [1, 2], "jax", [[1, 0], [1, 2]]),
     ],
 )
 @pytest.mark.jax
-def test_set_unitary_matrix(unitary_matrix, index, value, like, expected_matrix):
+def test_set_unitary_matrix(jax, unitary_matrix, index, value, like, expected_matrix):
     """Test the _set_unitary function on different interfaces."""
+    import jax.numpy as jnp
+
+    if jax:
+        unitary_matrix = jnp.array(unitary_matrix)
+        expected_matrix = jnp.array(expected_matrix)
 
     new_unitary_matrix = _set_unitary_matrix(unitary_matrix, index, value, like)
     assert qml.math.allclose(new_unitary_matrix, expected_matrix)

From 69acdf3d3ffbeb9b98c45dd96042a9891f547660 Mon Sep 17 00:00:00 2001
From: obliviateandsurrender <utkarshazad98@gmail.com>
Date: Fri, 18 Oct 2024 10:42:12 -0400
Subject: [PATCH 55/55] happy `pylint`

---
 tests/qchem/test_givens_rotations.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/qchem/test_givens_rotations.py b/tests/qchem/test_givens_rotations.py
index b28b78c9143..c7e76dbe391 100644
--- a/tests/qchem/test_givens_rotations.py
+++ b/tests/qchem/test_givens_rotations.py
@@ -145,6 +145,7 @@ def test_givens_matrix_exceptions():
         _givens_matrix(a, b)
 
 
+# pylint:disable = too-many-arguments
 @pytest.mark.parametrize(
     ("jax", "unitary_matrix", "index", "value", "like", "expected_matrix"),
     [