From f9cfdfa2668ff6091bc072e1877b764024a8e5cd Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Tue, 18 Jun 2024 16:00:36 -0400
Subject: [PATCH 01/18] Register apply_phaseshift and special cases.

---
 pennylane/devices/qubit/apply_operation.py | 40 ++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 4d7bbbe8038..31b3698d94b 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -406,6 +406,46 @@ def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None
     return math.stack([state[sl_0], state1], axis=axis)
 
 
+@apply_operation.register
+def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False, debugger=None, **_):
+    """Apply PhaseShift to state."""
+
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
+
+    # must be first state and then -1 because it breaks otherwise
+    state1 = math.multiply(state[sl_1], math.exp(1j * op.parameters[0]))
+    return math.stack([state[sl_0], state1], axis=axis)
+
+
+@apply_operation.register
+def apply_T(op: qml.T, state, is_state_batched: bool = False, debugger=None, **_):
+    """Apply T to state."""
+    return apply_operation(
+        qml.PhaseShift(np.pi / 4, op.wires),
+        state,
+        is_state_batched=is_state_batched,
+        debugger=debugger,
+    )
+
+
+@apply_operation.register
+def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_):
+    """Apply S to state."""
+    return apply_operation(
+        qml.PhaseShift(np.pi / 2, op.wires),
+        state,
+        is_state_batched=is_state_batched,
+        debugger=debugger,
+    )
+
+
 @apply_operation.register
 def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply cnot gate to state."""

From 3c43b52c8db4cd49e8a10e7f5f3d0665a2a4686f Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Tue, 18 Jun 2024 16:59:36 -0400
Subject: [PATCH 02/18] Add apply_controlled_operation

---
 pennylane/devices/qubit/apply_operation.py | 30 ++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 31b3698d94b..51e80b9bbd7 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -14,6 +14,7 @@
 """Functions to apply an operation to a state vector."""
 # pylint: disable=unused-argument, too-many-arguments
 
+import copy
 from functools import singledispatch
 from string import ascii_letters as alphabet
 
@@ -228,6 +229,35 @@ def _apply_operation_default(op, state, is_state_batched, debugger):
     return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
 
+@apply_operation.register
+def apply_controlled_operation(
+    op: qml.ops.op_math.ControlledOp, state, is_state_batched: bool = False, debugger=None, **_
+):  # pylint : disable=protected-access
+    if any(qml.math.get_deep_interface(data) != "numpy" for data in (op.data, state)):
+        return _apply_operation_default(
+            op, state, is_state_batched=is_state_batched, debugger=debugger
+        )
+    base = copy.deepcopy(op.base)
+    slices = [slice(None)] * qml.math.ndim(state)
+    for v, w in zip(op.control_values, op.control_wires):
+        slices[w] = int(v)
+    wires = list(range(qml.math.ndim(state)))
+    for w in reversed(sorted(op.control_wires)):
+        wires.pop(w)
+    op_wires = []
+    for b in base.wires:
+        for i, w in enumerate(wires):
+            if b == w:
+                op_wires.append(i)
+                break
+    base._wires = qml.wires.Wires(op_wires)
+    state = state + 0j
+    slc = state[tuple(slices)]
+    slc = apply_operation(base, slc, is_state_batched=is_state_batched, debugger=debugger)
+    state[tuple(slices)] = slc
+    return state
+
+
 @apply_operation.register
 def apply_conditional(
     op: Conditional,

From c873db2d7d3e285ce1c333ac748754423c0c9954 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Wed, 19 Jun 2024 13:54:08 -0400
Subject: [PATCH 03/18] Remove apply_controlled_operation

---
 pennylane/devices/qubit/apply_operation.py | 49 +++-------------------
 1 file changed, 6 insertions(+), 43 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 51e80b9bbd7..8ca4da3269c 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -14,7 +14,6 @@
 """Functions to apply an operation to a state vector."""
 # pylint: disable=unused-argument, too-many-arguments
 
-import copy
 from functools import singledispatch
 from string import ascii_letters as alphabet
 
@@ -229,35 +228,6 @@ def _apply_operation_default(op, state, is_state_batched, debugger):
     return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
 
-@apply_operation.register
-def apply_controlled_operation(
-    op: qml.ops.op_math.ControlledOp, state, is_state_batched: bool = False, debugger=None, **_
-):  # pylint : disable=protected-access
-    if any(qml.math.get_deep_interface(data) != "numpy" for data in (op.data, state)):
-        return _apply_operation_default(
-            op, state, is_state_batched=is_state_batched, debugger=debugger
-        )
-    base = copy.deepcopy(op.base)
-    slices = [slice(None)] * qml.math.ndim(state)
-    for v, w in zip(op.control_values, op.control_wires):
-        slices[w] = int(v)
-    wires = list(range(qml.math.ndim(state)))
-    for w in reversed(sorted(op.control_wires)):
-        wires.pop(w)
-    op_wires = []
-    for b in base.wires:
-        for i, w in enumerate(wires):
-            if b == w:
-                op_wires.append(i)
-                break
-    base._wires = qml.wires.Wires(op_wires)
-    state = state + 0j
-    slc = state[tuple(slices)]
-    slc = apply_operation(base, slc, is_state_batched=is_state_batched, debugger=debugger)
-    state[tuple(slices)] = slc
-    return state
-
-
 @apply_operation.register
 def apply_conditional(
     op: Conditional,
@@ -421,19 +391,12 @@ def apply_paulix(op: qml.X, state, is_state_batched: bool = False, debugger=None
 @apply_operation.register
 def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply pauliz to state."""
-
-    axis = op.wires[0] + is_state_batched
-    n_dim = math.ndim(state)
-
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
-
-    sl_0 = _get_slice(0, axis, n_dim)
-    sl_1 = _get_slice(1, axis, n_dim)
-
-    # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(state[sl_1], -1)
-    return math.stack([state[sl_0], state1], axis=axis)
+    return apply_operation(
+        qml.PhaseShift(np.pi, op.wires),
+        state,
+        is_state_batched=is_state_batched,
+        debugger=debugger,
+    )
 
 
 @apply_operation.register

From 2d9bfe4a53077ef110c1b96894af55069aba3048 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 13:19:58 +0000
Subject: [PATCH 04/18] Deal with PhaseShift batched.

---
 pennylane/devices/qubit/apply_operation.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 8ca4da3269c..6526bc22771 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -403,6 +403,19 @@ def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None
 def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply PhaseShift to state."""
 
+    params = op.parameters[0]
+    if is_state_batched or (op.batch_size is not None and len(params) > 1):
+        slices = []
+        for i, p in enumerate(op.parameters[0]):
+            slices.append(
+                apply_phaseshift(
+                    qml.PhaseShift(p, wires=op.wires),
+                    state[i] if is_state_batched else state,
+                    is_state_batched=False,
+                )
+            )
+        return math.stack(slices, axis=0)
+
     axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
@@ -414,7 +427,10 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
 
     # must be first state and then -1 because it breaks otherwise
     state1 = math.multiply(state[sl_1], math.exp(1j * op.parameters[0]))
-    return math.stack([state[sl_0], state1], axis=axis)
+    state = math.stack([state[sl_0], state1], axis=axis)
+    if op.batch_size == 1:
+        state = math.stack([state], axis=0)
+    return state
 
 
 @apply_operation.register

From ea4e2560db1badefd8a770804a45dd5b99d95c4f Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 13:41:24 +0000
Subject: [PATCH 05/18] Fix autograd bug.

---
 pennylane/devices/qubit/apply_operation.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 6526bc22771..1fd21dfaee4 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -403,10 +403,16 @@ def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None
 def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply PhaseShift to state."""
 
-    params = op.parameters[0]
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    params = qml.math.atleast_1d(op.parameters[0])
     if is_state_batched or (op.batch_size is not None and len(params) > 1):
         slices = []
-        for i, p in enumerate(op.parameters[0]):
+        for i, p in enumerate(params):
             slices.append(
                 apply_phaseshift(
                     qml.PhaseShift(p, wires=op.wires),
@@ -416,17 +422,11 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
             )
         return math.stack(slices, axis=0)
 
-    axis = op.wires[0] + is_state_batched
-    n_dim = math.ndim(state)
-
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
-
     sl_0 = _get_slice(0, axis, n_dim)
     sl_1 = _get_slice(1, axis, n_dim)
 
     # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(state[sl_1], math.exp(1j * op.parameters[0]))
+    state1 = math.cast(state[sl_1], dtype=complex) * math.exp(1.0j * params)
     state = math.stack([state[sl_0], state1], axis=axis)
     if op.batch_size == 1:
         state = math.stack([state], axis=0)

From ee5542b1f513a51c1f019177492be745688aee2f Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 16:08:15 +0000
Subject: [PATCH 06/18] Fix TF tests.

---
 pennylane/devices/qubit/apply_operation.py  | 17 +++++++++--------
 tests/devices/qubit/test_apply_operation.py |  4 ++--
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 1fd21dfaee4..c64725f9e9b 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -403,14 +403,9 @@ def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None
 def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply PhaseShift to state."""
 
-    axis = op.wires[0] + is_state_batched
-    n_dim = math.ndim(state)
-
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
-
-    params = qml.math.atleast_1d(op.parameters[0])
+    params = op.parameters[0]
     if is_state_batched or (op.batch_size is not None and len(params) > 1):
+        params = math.atleast_1d(params)
         slices = []
         for i, p in enumerate(params):
             slices.append(
@@ -422,11 +417,17 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
             )
         return math.stack(slices, axis=0)
 
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
     sl_0 = _get_slice(0, axis, n_dim)
     sl_1 = _get_slice(1, axis, n_dim)
 
     # must be first state and then -1 because it breaks otherwise
-    state1 = math.cast(state[sl_1], dtype=complex) * math.exp(1.0j * params)
+    state1 = math.multiply(state[sl_1], math.exp(1j * params))
     state = math.stack([state[sl_0], state1], axis=axis)
     if op.batch_size == 1:
         state = math.stack([state], axis=0)
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 3e517879b78..5d6f1da11bf 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -802,7 +802,7 @@ def test_broadcasted_op(self, op, method, ml_framework):
     @pytest.mark.parametrize("op", unbroadcasted_ops)
     def test_broadcasted_state(self, op, method, ml_framework):
         """Tests that unbatched operations are applied correctly to a batched state."""
-        state = np.ones((3, 2, 2, 2)) / np.sqrt(8)
+        state = np.ones((3, 2, 2, 2), dtype=complex) / np.sqrt(8)
 
         res = method(op, qml.math.asarray(state, like=ml_framework), is_state_batched=True)
         missing_wires = 3 - len(op.wires)
@@ -819,7 +819,7 @@ def test_broadcasted_op_broadcasted_state(self, op, method, ml_framework):
         if method is apply_operation_tensordot:
             pytest.skip("Tensordot doesn't support batched operator and batched state.")
 
-        state = np.ones((3, 2, 2, 2)) / np.sqrt(8)
+        state = np.ones((3, 2, 2, 2), dtype=complex) / np.sqrt(8)
 
         res = method(op, qml.math.asarray(state, like=ml_framework), is_state_batched=True)
         missing_wires = 3 - len(op.wires)

From 80505c6ee21c511b346b68ab575c53d053a50906 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 19:14:44 +0000
Subject: [PATCH 07/18] Do not reuse PhaseShift.

---
 pennylane/devices/qubit/apply_operation.py | 71 ++++++++++++++--------
 1 file changed, 46 insertions(+), 25 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 0de51992c87..31201dac679 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -378,8 +378,8 @@ def apply_identity(op: qml.Identity, state, is_state_batched: bool = False, debu
 def apply_global_phase(
     op: qml.GlobalPhase, state, is_state_batched: bool = False, debugger=None, **_
 ):
-    """Applies a :class:`~.GlobalPhase` operation by multiplying the state by ``exp(1j * op.data[0])``"""
-    return qml.math.exp(-1j * qml.math.cast(op.data[0], complex)) * state
+    """Applies a :class:`~.GlobalPhase` operation by multiplying the state by ``exp(1.0j * op.data[0])``"""
+    return qml.math.exp(-1.0j * qml.math.cast(op.data[0], complex)) * state
 
 
 @apply_operation.register
@@ -392,12 +392,19 @@ def apply_paulix(op: qml.X, state, is_state_batched: bool = False, debugger=None
 @apply_operation.register
 def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply pauliz to state."""
-    return apply_operation(
-        qml.PhaseShift(np.pi, op.wires),
-        state,
-        is_state_batched=is_state_batched,
-        debugger=debugger,
-    )
+
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
+
+    # must be first state and then -1 because it breaks otherwise
+    state1 = math.multiply(state[sl_1], -1)
+    return math.stack([state[sl_0], state1], axis=axis)
 
 
 @apply_operation.register
@@ -428,7 +435,7 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     sl_1 = _get_slice(1, axis, n_dim)
 
     # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(state[sl_1], math.exp(1j * params))
+    state1 = math.multiply(state[sl_1], math.exp(1.0j * params))
     state = math.stack([state[sl_0], state1], axis=axis)
     if op.batch_size == 1:
         state = math.stack([state], axis=0)
@@ -438,23 +445,37 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
 @apply_operation.register
 def apply_T(op: qml.T, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply T to state."""
-    return apply_operation(
-        qml.PhaseShift(np.pi / 4, op.wires),
-        state,
-        is_state_batched=is_state_batched,
-        debugger=debugger,
-    )
 
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
 
-@apply_operation.register
-def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_):
-    """Apply S to state."""
-    return apply_operation(
-        qml.PhaseShift(np.pi / 2, op.wires),
-        state,
-        is_state_batched=is_state_batched,
-        debugger=debugger,
-    )
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
+
+    # must be first state and then -1 because it breaks otherwise
+    state1 = math.multiply(state[sl_1], math.exp(0.25j * np.pi))
+    return math.stack([state[sl_0], state1], axis=axis)
+
+
+# @apply_operation.register
+# def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_):
+#     """Apply T to state."""
+
+#     axis = op.wires[0] + is_state_batched
+#     n_dim = math.ndim(state)
+
+#     if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+#         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+#     sl_0 = _get_slice(0, axis, n_dim)
+#     sl_1 = _get_slice(1, axis, n_dim)
+
+#     # must be first state and then -1 because it breaks otherwise
+#     state1 = math.multiply(state[sl_1], 1.0j)
+#     return math.stack([state[sl_0], state1], axis=axis)
 
 
 @apply_operation.register
@@ -687,7 +708,7 @@ def _evolve_state_vector_under_parametrized_evolution(
 
     def fun(y, t):
         """dy/dt = -i H(t) y"""
-        return (-1j * H_jax(operation.data, t=t)) @ y
+        return (-1.0j * H_jax(operation.data, t=t)) @ y
 
     result = odeint(fun, state, operation.t, **operation.odeint_kwargs)
     if operation.hyperparameters["return_intermediate"]:

From 6358bba0435b1354a83f7fd125e435bec5f33f07 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 19:24:32 +0000
Subject: [PATCH 08/18] Put back S.

---
 pennylane/devices/qubit/apply_operation.py | 34 +++++++++++-----------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 31201dac679..d6fcf5e1ee3 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -378,8 +378,8 @@ def apply_identity(op: qml.Identity, state, is_state_batched: bool = False, debu
 def apply_global_phase(
     op: qml.GlobalPhase, state, is_state_batched: bool = False, debugger=None, **_
 ):
-    """Applies a :class:`~.GlobalPhase` operation by multiplying the state by ``exp(1.0j * op.data[0])``"""
-    return qml.math.exp(-1.0j * qml.math.cast(op.data[0], complex)) * state
+    """Applies a :class:`~.GlobalPhase` operation by multiplying the state by ``exp(1j * op.data[0])``"""
+    return qml.math.exp(-1j * qml.math.cast(op.data[0], complex)) * state
 
 
 @apply_operation.register
@@ -435,7 +435,7 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     sl_1 = _get_slice(1, axis, n_dim)
 
     # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(state[sl_1], math.exp(1.0j * params))
+    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), math.exp(1j * params))
     state = math.stack([state[sl_0], state1], axis=axis)
     if op.batch_size == 1:
         state = math.stack([state], axis=0)
@@ -456,26 +456,26 @@ def apply_T(op: qml.T, state, is_state_batched: bool = False, debugger=None, **_
     sl_1 = _get_slice(1, axis, n_dim)
 
     # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(state[sl_1], math.exp(0.25j * np.pi))
+    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), math.exp(0.25j * np.pi))
     return math.stack([state[sl_0], state1], axis=axis)
 
 
-# @apply_operation.register
-# def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_):
-#     """Apply T to state."""
+@apply_operation.register
+def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_):
+    """Apply S to state."""
 
-#     axis = op.wires[0] + is_state_batched
-#     n_dim = math.ndim(state)
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
 
-#     if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-#         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
-#     sl_0 = _get_slice(0, axis, n_dim)
-#     sl_1 = _get_slice(1, axis, n_dim)
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
 
-#     # must be first state and then -1 because it breaks otherwise
-#     state1 = math.multiply(state[sl_1], 1.0j)
-#     return math.stack([state[sl_0], state1], axis=axis)
+    # must be first state and then -1 because it breaks otherwise
+    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), 1j)
+    return math.stack([state[sl_0], state1], axis=axis)
 
 
 @apply_operation.register
@@ -708,7 +708,7 @@ def _evolve_state_vector_under_parametrized_evolution(
 
     def fun(y, t):
         """dy/dt = -i H(t) y"""
-        return (-1.0j * H_jax(operation.data, t=t)) @ y
+        return (-1j * H_jax(operation.data, t=t)) @ y
 
     result = odeint(fun, state, operation.t, **operation.odeint_kwargs)
     if operation.hyperparameters["return_intermediate"]:

From aaa32ff04c0e86e69e4a4330389978ead49a9461 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 19:30:49 +0000
Subject: [PATCH 09/18] Fix complex issue with TF

---
 pennylane/devices/qubit/apply_operation.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index d6fcf5e1ee3..aa1bf83f0f2 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -435,7 +435,9 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     sl_1 = _get_slice(1, axis, n_dim)
 
     # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), math.exp(1j * params))
+    state1 = math.multiply(
+        math.cast(state[sl_1], dtype=complex), math.exp(1j * math.cast(params, dtype=complex))
+    )
     state = math.stack([state[sl_0], state1], axis=axis)
     if op.batch_size == 1:
         state = math.stack([state], axis=0)

From 0e3700e2557b78a6792b62fc41abbde63fcc3039 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 20:04:57 +0000
Subject: [PATCH 10/18] Use PhaseShift in T, S, Z.

---
 pennylane/devices/qubit/apply_operation.py | 57 +++++++---------------
 1 file changed, 18 insertions(+), 39 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index aa1bf83f0f2..8e8a34af660 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -392,19 +392,12 @@ def apply_paulix(op: qml.X, state, is_state_batched: bool = False, debugger=None
 @apply_operation.register
 def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply pauliz to state."""
-
-    axis = op.wires[0] + is_state_batched
-    n_dim = math.ndim(state)
-
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
-
-    sl_0 = _get_slice(0, axis, n_dim)
-    sl_1 = _get_slice(1, axis, n_dim)
-
-    # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(state[sl_1], -1)
-    return math.stack([state[sl_0], state1], axis=axis)
+    return apply_operation(
+        qml.PhaseShift(np.pi, op.wires),
+        state,
+        is_state_batched=is_state_batched,
+        debugger=debugger,
+    )
 
 
 @apply_operation.register
@@ -447,37 +440,23 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
 @apply_operation.register
 def apply_T(op: qml.T, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply T to state."""
-
-    axis = op.wires[0] + is_state_batched
-    n_dim = math.ndim(state)
-
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
-
-    sl_0 = _get_slice(0, axis, n_dim)
-    sl_1 = _get_slice(1, axis, n_dim)
-
-    # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), math.exp(0.25j * np.pi))
-    return math.stack([state[sl_0], state1], axis=axis)
+    return apply_operation(
+        qml.PhaseShift(np.pi / 4, op.wires),
+        state,
+        is_state_batched=is_state_batched,
+        debugger=debugger,
+    )
 
 
 @apply_operation.register
 def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply S to state."""
-
-    axis = op.wires[0] + is_state_batched
-    n_dim = math.ndim(state)
-
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
-
-    sl_0 = _get_slice(0, axis, n_dim)
-    sl_1 = _get_slice(1, axis, n_dim)
-
-    # must be first state and then -1 because it breaks otherwise
-    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), 1j)
-    return math.stack([state[sl_0], state1], axis=axis)
+    return apply_operation(
+        qml.PhaseShift(np.pi / 2, op.wires),
+        state,
+        is_state_batched=is_state_batched,
+        debugger=debugger,
+    )
 
 
 @apply_operation.register

From dfdc39bef98073128d4b96ec8f205379fe4215bd Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 20:16:01 +0000
Subject: [PATCH 11/18] Update changelog

---
 doc/releases/changelog-dev.md | 29 ++++++++++++++++-------------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index eb52d55607d..6caaae84c7e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,10 +4,10 @@
 
 <h3>New features since last release</h3>
 
-* Added a quantum debugger (`PLDB`) which interfaces via `qml.breakpoint()` and provides tools for 
+* Added a quantum debugger (`PLDB`) which interfaces via `qml.breakpoint()` and provides tools for
   debugging quantum circuits. Users can step through the quantum circuit operations, dynamically
-  queue operations and make measurements using (`qml.debug_state()`, `qml.debug_probs()`, 
-  `qml.debug_expval()`, and `qml.debug_tape()`). Consider the following python script 
+  queue operations and make measurements using (`qml.debug_state()`, `qml.debug_probs()`,
+  `qml.debug_expval()`, and `qml.debug_tape()`). Consider the following python script
   containing the quantum circuit with breakpoints.
   [(#5680)](https://github.com/PennyLaneAI/pennylane/pull/5680)
   [(#5749)](https://github.com/PennyLaneAI/pennylane/pull/5749)
@@ -56,7 +56,7 @@
 * The `default.tensor` device is introduced to perform tensor network simulations of quantum circuits using the `mps` (Matrix Product State) method.
   [(#5699)](https://github.com/PennyLaneAI/pennylane/pull/5699)
 
-* A new `qml.noise` module which contains utility function for building `NoiseModels` 
+* A new `qml.noise` module which contains utility function for building `NoiseModels`
   and an `add_noise` tranform for addding it to quantum circuits.
   [(#5674)](https://github.com/PennyLaneAI/pennylane/pull/5674)
   [(#5684)](https://github.com/PennyLaneAI/pennylane/pull/5684)
@@ -93,7 +93,7 @@
   ───PhaseDamping(0.40)──────────────────────────┤ ╰<Z@Z>
   ```
 
-* The ``from_openfermion`` and ``to_openfermion`` functions are added to convert between 
+* The ``from_openfermion`` and ``to_openfermion`` functions are added to convert between
   OpenFermion and PennyLane objects.
   [(#5773)](https://github.com/PennyLaneAI/pennylane/pull/5773)
   [(#5808)](https://github.com/PennyLaneAI/pennylane/pull/5808)
@@ -113,16 +113,19 @@
 
 <h3>Improvements 🛠</h3>
 
+* Port the fast `apply_operation` implementation of `PauliZ` to `PhaseShift`, `S` and `T`.
+  [(#5876)](https://github.com/PennyLaneAI/pennylane/pull/5876)
+
 * Add operation and measurement specific routines in `default.tensor` to improve scalability.
   [(#5795)](https://github.com/PennyLaneAI/pennylane/pull/5795)
-  
+
 * `param_shift` with the `broadcast=True` option now supports shot vectors and multiple measurements.
   [(#5667)](https://github.com/PennyLaneAI/pennylane/pull/5667)
 
 * `default.clifford` now supports arbitrary state-based measurements with `qml.Snapshot`.
   [(#5794)](https://github.com/PennyLaneAI/pennylane/pull/5794)
 
-* `qml.TrotterProduct` is now compatible with resource tracking by inheriting from `ResourcesOperation`. 
+* `qml.TrotterProduct` is now compatible with resource tracking by inheriting from `ResourcesOperation`.
    [(#5680)](https://github.com/PennyLaneAI/pennylane/pull/5680)
 
 * The wires for the `default.tensor` device are selected at runtime if they are not provided by user.
@@ -191,7 +194,7 @@
 
 * The `dynamic_one_shot` transform is made compatible with the Catalyst compiler.
   [(#5766)](https://github.com/PennyLaneAI/pennylane/pull/5766)
-  
+
 * Rationalize MCM tests, removing most end-to-end tests from the native MCM test file,
   but keeping one that validates multiple mid-circuit measurements with any allowed return
   and interface end-to-end tests.
@@ -322,7 +325,7 @@
 * The qchem docs are updated with the new qchem improvements.
   [(#5758)](https://github.com/PennyLaneAI/pennylane/pull/5758/)
   [(#5638)](https://github.com/PennyLaneAI/pennylane/pull/5638/)
-  
+
 * `specs()` can now be requested at any specific point of the transform program through the `level` keyword argument.
   [(#5781)](https://github.com/PennyLaneAI/pennylane/pull/5781/)
 
@@ -347,10 +350,10 @@
 
 * Implemented kwargs (`check_interface`, `check_trainability`, `rtol` and `atol`) support in `qml.equal` for the operators `Pow`, `Adjoint`, `Exp`, and `SProd`.
   [(#5668)](https://github.com/PennyLaneAI/pennylane/issues/5668)
-  
+
 * `qml.QutritDepolarizingChannel` has been added, allowing for depolarizing noise to be simulated on the `default.qutrit.mixed` device.
   [(#5502)](https://github.com/PennyLaneAI/pennylane/pull/5502)
- 
+
 * Implement support in `assert_equal` for `Operator`, `Controlled`, `Adjoint`, `Pow`, `Exp`, `SProd`, `ControlledSequence`, `Prod`, `Sum`, `Tensor` and `Hamiltonian`
   [(#5780)](https://github.com/PennyLaneAI/pennylane/pull/5780)
   [(#5877)](https://github.com/PennyLaneAI/pennylane/pull/5877)
@@ -362,8 +365,8 @@
   [(#5503)](https://github.com/PennyLaneAI/pennylane/pull/5503)
   [(#5757)](https://github.com/PennyLaneAI/pennylane/pull/5757)
   [(#5799)](https://github.com/PennyLaneAI/pennylane/pull/5799)
-  
-* `qml.TritFlip` has been added, allowing for trit flip errors, such as misclassification, 
+
+* `qml.TritFlip` has been added, allowing for trit flip errors, such as misclassification,
   to be simulated on the `default.qutrit.mixed` device.
   [(#5784)](https://github.com/PennyLaneAI/pennylane/pull/5784)
 

From 1ebb499bc9d2f40d0b34c05cf8014c6caf5a8073 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 20 Jun 2024 20:33:17 +0000
Subject: [PATCH 12/18] Revert to full implementations.

---
 pennylane/devices/qubit/apply_operation.py | 74 ++++++++++++++++------
 1 file changed, 55 insertions(+), 19 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 8e8a34af660..602dbc0fa91 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -389,15 +389,40 @@ def apply_paulix(op: qml.X, state, is_state_batched: bool = False, debugger=None
     return math.roll(state, 1, axis)
 
 
+@apply_operation.register
+def apply_pauliy(op: qml.Y, state, is_state_batched: bool = False, debugger=None, **_):
+    """Apply pauliy to state."""
+
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
+
+    state0 = math.multiply(math.cast(state[sl_0], dtype=complex), 1j)
+    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), -1j)
+    return math.stack([state1, state0], axis=axis)
+
+
 @apply_operation.register
 def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply pauliz to state."""
-    return apply_operation(
-        qml.PhaseShift(np.pi, op.wires),
-        state,
-        is_state_batched=is_state_batched,
-        debugger=debugger,
-    )
+
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
+
+    # must be first state and then -1 because it breaks otherwise
+    state1 = math.multiply(state[sl_1], -1)
+    return math.stack([state[sl_0], state1], axis=axis)
 
 
 @apply_operation.register
@@ -427,7 +452,6 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     sl_0 = _get_slice(0, axis, n_dim)
     sl_1 = _get_slice(1, axis, n_dim)
 
-    # must be first state and then -1 because it breaks otherwise
     state1 = math.multiply(
         math.cast(state[sl_1], dtype=complex), math.exp(1j * math.cast(params, dtype=complex))
     )
@@ -440,23 +464,35 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
 @apply_operation.register
 def apply_T(op: qml.T, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply T to state."""
-    return apply_operation(
-        qml.PhaseShift(np.pi / 4, op.wires),
-        state,
-        is_state_batched=is_state_batched,
-        debugger=debugger,
-    )
+
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
+
+    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), math.exp(0.25j * np.pi))
+    return math.stack([state[sl_0], state1], axis=axis)
 
 
 @apply_operation.register
 def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply S to state."""
-    return apply_operation(
-        qml.PhaseShift(np.pi / 2, op.wires),
-        state,
-        is_state_batched=is_state_batched,
-        debugger=debugger,
-    )
+
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
+
+    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), 1j)
+    return math.stack([state[sl_0], state1], axis=axis)
 
 
 @apply_operation.register

From 60ac35b7c467ffb709a0a8ded776f0166cd7937e Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Tue, 25 Jun 2024 12:09:35 +0000
Subject: [PATCH 13/18] pragma: no cover

---
 pennylane/devices/qubit/apply_operation.py | 26 ++++------------------
 1 file changed, 4 insertions(+), 22 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 602dbc0fa91..cfc84a00c18 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -389,24 +389,6 @@ def apply_paulix(op: qml.X, state, is_state_batched: bool = False, debugger=None
     return math.roll(state, 1, axis)
 
 
-@apply_operation.register
-def apply_pauliy(op: qml.Y, state, is_state_batched: bool = False, debugger=None, **_):
-    """Apply pauliy to state."""
-
-    axis = op.wires[0] + is_state_batched
-    n_dim = math.ndim(state)
-
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
-
-    sl_0 = _get_slice(0, axis, n_dim)
-    sl_1 = _get_slice(1, axis, n_dim)
-
-    state0 = math.multiply(math.cast(state[sl_0], dtype=complex), 1j)
-    state1 = math.multiply(math.cast(state[sl_1], dtype=complex), -1j)
-    return math.stack([state1, state0], axis=axis)
-
-
 @apply_operation.register
 def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply pauliz to state."""
@@ -446,7 +428,7 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, axis, n_dim)
@@ -468,7 +450,7 @@ def apply_T(op: qml.T, state, is_state_batched: bool = False, debugger=None, **_
     axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, axis, n_dim)
@@ -485,7 +467,7 @@ def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_
     axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, axis, n_dim)
@@ -502,7 +484,7 @@ def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=Non
     control_axes = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, control_axes, n_dim)

From fa0cc2857871861e1e3ba5a4cf1cae02c4bdebf5 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Mon, 8 Jul 2024 16:44:50 +0000
Subject: [PATCH 14/18] Broadcast params in phaseshift.

---
 pennylane/devices/qubit/apply_operation.py | 36 ++++++++++------------
 1 file changed, 16 insertions(+), 20 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index cfc84a00c18..5648dcd0e6e 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -411,34 +411,30 @@ def apply_pauliz(op: qml.Z, state, is_state_batched: bool = False, debugger=None
 def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply PhaseShift to state."""
 
-    params = op.parameters[0]
-    if is_state_batched or (op.batch_size is not None and len(params) > 1):
-        params = math.atleast_1d(params)
-        slices = []
-        for i, p in enumerate(params):
-            slices.append(
-                apply_phaseshift(
-                    qml.PhaseShift(p, wires=op.wires),
-                    state[i] if is_state_batched else state,
-                    is_state_batched=False,
-                )
-            )
-        return math.stack(slices, axis=0)
-
-    axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
     if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
+    axis = op.wires[0] + is_state_batched
+
     sl_0 = _get_slice(0, axis, n_dim)
     sl_1 = _get_slice(1, axis, n_dim)
 
-    state1 = math.multiply(
-        math.cast(state[sl_1], dtype=complex), math.exp(1j * math.cast(params, dtype=complex))
-    )
-    state = math.stack([state[sl_0], state1], axis=axis)
-    if op.batch_size == 1:
+    params = math.cast(op.parameters[0], dtype=complex)
+    state0 = state[sl_0]
+    state1 = state[sl_1]
+    if not is_state_batched and params.size > 1:
+        params = qml.math.expand_dims(params, tuple(range(1, n_dim)))
+        state0 = qml.math.expand_dims(state0, 0)
+        state0 = qml.math.repeat(state0, params.size, axis=0)
+        state1 = qml.math.expand_dims(state1, 0)
+        axis = axis + 1
+    elif params.size > 1:
+        params = qml.math.expand_dims(params, tuple(range(1, n_dim - 1)))
+    state1 = math.multiply(math.cast(state1, dtype=complex), math.exp(1j * params))
+    state = math.stack([state0, state1], axis=axis)
+    if not is_state_batched and op.batch_size == 1:
         state = math.stack([state], axis=0)
     return state
 

From e4e789877d140118979e293a5330939cdad7b56d Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Mon, 8 Jul 2024 18:05:03 +0000
Subject: [PATCH 15/18] Fix broadcasting for all interfaces.

---
 pennylane/devices/qubit/apply_operation.py  | 19 ++++++++++---------
 tests/devices/qubit/test_apply_operation.py |  2 +-
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 5648dcd0e6e..bd74984c7ed 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -424,15 +424,16 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     params = math.cast(op.parameters[0], dtype=complex)
     state0 = state[sl_0]
     state1 = state[sl_1]
-    if not is_state_batched and params.size > 1:
-        params = qml.math.expand_dims(params, tuple(range(1, n_dim)))
-        state0 = qml.math.expand_dims(state0, 0)
-        state0 = qml.math.repeat(state0, params.size, axis=0)
-        state1 = qml.math.expand_dims(state1, 0)
-        axis = axis + 1
-    elif params.size > 1:
-        params = qml.math.expand_dims(params, tuple(range(1, n_dim - 1)))
-    state1 = math.multiply(math.cast(state1, dtype=complex), math.exp(1j * params))
+    if op.batch_size is not None and len(params) > 1:
+        params = math.array(params, like=math.get_deep_interface(state))
+        if is_state_batched:
+            params = qml.math.reshape(params, (-1,) + (1,) * (n_dim - 2))
+        else:
+            axis = axis + 1
+            params = qml.math.reshape(params, (-1,) + (1,) * (n_dim - 1))
+            state0 = qml.math.expand_dims(state0, 0) + math.zeros_like(params)
+            state1 = qml.math.expand_dims(state1, 0)
+    state1 = math.multiply(math.cast(state1, dtype=complex), math.exp(1.0j * params))
     state = math.stack([state0, state1], axis=axis)
     if not is_state_batched and op.batch_size == 1:
         state = math.stack([state], axis=0)
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 282b7f4b12d..f4246d525b3 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -778,7 +778,7 @@ class TestBroadcasting:  # pylint: disable=too-few-public-methods
     @pytest.mark.parametrize("op", broadcasted_ops)
     def test_broadcasted_op(self, op, method, ml_framework):
         """Tests that batched operations are applied correctly to an unbatched state."""
-        state = np.ones((2, 2, 2)) / np.sqrt(8)
+        state = np.ones((2, 2, 2), dtype=complex) / np.sqrt(8)
 
         res = method(op, qml.math.asarray(state, like=ml_framework))
         missing_wires = 3 - len(op.wires)

From 16897223000440391c053a2f45fb4e555a1e9472 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Mon, 8 Jul 2024 18:37:25 +0000
Subject: [PATCH 16/18] Fix jax issue with math.array.

---
 pennylane/devices/qubit/apply_operation.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index bd74984c7ed..3b1214a36f2 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -425,7 +425,9 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     state0 = state[sl_0]
     state1 = state[sl_1]
     if op.batch_size is not None and len(params) > 1:
-        params = math.array(params, like=math.get_deep_interface(state))
+        interface = math.get_deep_interface(state)
+        if interface == "torch":
+            params = math.array(params, like=interface)
         if is_state_batched:
             params = qml.math.reshape(params, (-1,) + (1,) * (n_dim - 2))
         else:

From b309faf0fcfb3239d4e822ce20df0c23fe961bce Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Mon, 8 Jul 2024 20:21:44 +0000
Subject: [PATCH 17/18] Add tests and remove obsolete pragma: no cover

---
 pennylane/devices/qubit/apply_operation.py  | 10 +++++-----
 tests/devices/qubit/test_apply_operation.py |  6 ++++--
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 3b1214a36f2..f339befedc6 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -413,7 +413,7 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
 
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     axis = op.wires[0] + is_state_batched
@@ -425,7 +425,7 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
     state0 = state[sl_0]
     state1 = state[sl_1]
     if op.batch_size is not None and len(params) > 1:
-        interface = math.get_deep_interface(state)
+        interface = math.get_interface(state)
         if interface == "torch":
             params = math.array(params, like=interface)
         if is_state_batched:
@@ -449,7 +449,7 @@ def apply_T(op: qml.T, state, is_state_batched: bool = False, debugger=None, **_
     axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, axis, n_dim)
@@ -466,7 +466,7 @@ def apply_S(op: qml.S, state, is_state_batched: bool = False, debugger=None, **_
     axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, axis, n_dim)
@@ -483,7 +483,7 @@ def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=Non
     control_axes = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
-    if n_dim >= 9 and math.get_interface(state) == "tensorflow":  # pragma: no cover
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
         return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, control_axes, n_dim)
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index f4246d525b3..fe52935b670 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -1226,13 +1226,15 @@ def test_with_torch(self, batch_dim):
 class TestLargeTFCornerCases:
     """Test large corner cases for tensorflow."""
 
-    @pytest.mark.parametrize("op", (qml.PauliZ(8), qml.CNOT((5, 6))))
+    @pytest.mark.parametrize(
+        "op", (qml.PauliZ(8), qml.PhaseShift(1.0, 8), qml.S(8), qml.T(8), qml.CNOT((5, 6)))
+    )
     def test_tf_large_state(self, op):
         """Tests that custom kernels that use slicing fall back to a different method when
         the state has a large number of wires."""
         import tensorflow as tf
 
-        state = np.zeros([2] * 10)
+        state = np.zeros([2] * 10, dtype=complex)
         state = tf.Variable(state)
         new_state = apply_operation(op, state)
 

From 75af4b6b0e8397cb45ef67a9939162656b6eeeb2 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Tue, 9 Jul 2024 13:31:07 -0400
Subject: [PATCH 18/18] Update pennylane/devices/qubit/apply_operation.py

Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
---
 pennylane/devices/qubit/apply_operation.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index f339befedc6..e51cad39d9e 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -429,12 +429,12 @@ def apply_phaseshift(op: qml.PhaseShift, state, is_state_batched: bool = False,
         if interface == "torch":
             params = math.array(params, like=interface)
         if is_state_batched:
-            params = qml.math.reshape(params, (-1,) + (1,) * (n_dim - 2))
+            params = math.reshape(params, (-1,) + (1,) * (n_dim - 2))
         else:
             axis = axis + 1
-            params = qml.math.reshape(params, (-1,) + (1,) * (n_dim - 1))
-            state0 = qml.math.expand_dims(state0, 0) + math.zeros_like(params)
-            state1 = qml.math.expand_dims(state1, 0)
+            params = math.reshape(params, (-1,) + (1,) * (n_dim - 1))
+            state0 = math.expand_dims(state0, 0) + math.zeros_like(params)
+            state1 = math.expand_dims(state1, 0)
     state1 = math.multiply(math.cast(state1, dtype=complex), math.exp(1.0j * params))
     state = math.stack([state0, state1], axis=axis)
     if not is_state_batched and op.batch_size == 1: