Optimize PhaseShift, T, S gates

doc/releases/changelog-dev.md

@@ -12,6 +12,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* Port the fast `apply_operation` implementation of `PauliZ` to `PhaseShift`, `S` and `T`.
+ [(#5876)](https://github.com/PennyLaneAI/pennylane/pull/5876)
+
 * `qml.UCCSD` now accepts an additional optional argument, `n_repeats`, which defines the number of
  times the UCCSD template is repeated. This can improve the accuracy of the template by reducing
  the Trotter error but would result in deeper circuits.
@@ -31,8 +34,8 @@
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
-
 Yushao Chen,
 Christina Lee,
 William Maxwell,
+Vincent Michaud-Rioux,
 Erik Schultheis.

pennylane/devices/qubit/apply_operation.py

@@ -407,14 +407,80 @@ 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."""
+
+ 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)
+
+ 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
+
+
+@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": # pragma: no cover
+ 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."""
+
+ 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)
+
+ 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
 def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=None, **_):
  """Apply cnot gate to state."""
  target_axes = (op.wires[1] - 1 if op.wires[1] > op.wires[0] else op.wires[1]) + is_state_batched
  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)

tests/devices/qubit/test_apply_operation.py

@@ -796,7 +796,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)
@@ -813,7 +813,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)