qml.Projector compatibility with new default qubit (#4452)

* compatibility with DefaultQubit2

* update changelog

* testing tests

* extensive testing

* more tests

* Fix `__reduce__` docstring

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* Fix test

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* simplify typing

* formatting

* simplify measurement logic

* serialize without

* remove copy and extend serialization tests

* remove redundant inheritance

---------

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

doc/releases/changelog-dev.md

@@ -310,6 +310,9 @@ array([False, False])
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.Projector` is pickle-able again.
+ [(#4452)](https://github.com/PennyLaneAI/pennylane/pull/4452)
+
 * Allow sparse matrix calculation of `SProd`s containing a `Tensor`. When using
  `Tensor.sparse_matrix()`, it is recommended to use the `wire_order` keyword argument over `wires`. 
  [(#4424)](https://github.com/PennyLaneAI/pennylane/pull/4424)

pennylane/_qubit_device.py

@@ -63,6 +63,7 @@
  VnEntropyMP,
  Shots,
 )
+from pennylane.ops.qubit.observables import BasisStateProjector
 from pennylane.resource import Resources
 from pennylane.operation import operation_derivative, Operation
 from pennylane.tape import QuantumScript, QuantumTape
@@ -1674,9 +1675,7 @@ def marginal_prob(self, prob, wires=None):
  return self._reshape(prob, flat_shape)
 
  def expval(self, observable, shot_range=None, bin_size=None):
- if observable.name == "Projector" and len(observable.parameters[0]) == len(
- observable.wires
- ):
+ if isinstance(observable, BasisStateProjector):
  # branch specifically to handle the basis state projector observable
  idx = int("".join(str(i) for i in observable.parameters[0]), 2)
  probs = self.probability(
@@ -1706,9 +1705,7 @@ def expval(self, observable, shot_range=None, bin_size=None):
  return np.squeeze(np.mean(samples, axis=axis))
 
  def var(self, observable, shot_range=None, bin_size=None):
- if observable.name == "Projector" and len(observable.parameters[0]) == len(
- observable.wires
- ):
+ if isinstance(observable, BasisStateProjector):
  # branch specifically to handle the basis state projector observable
  idx = int("".join(str(i) for i in observable.parameters[0]), 2)
  probs = self.probability(

pennylane/measurements/expval.py

@@ -19,7 +19,7 @@
 
 import pennylane as qml
 from pennylane.operation import Operator
-from pennylane.ops import Projector
+from pennylane.ops.qubit.observables import BasisStateProjector
 from pennylane.wires import Wires
 
 from .measurements import Expectation, SampleMeasurement, StateMeasurement
@@ -104,8 +104,8 @@ def process_samples(
  shot_range: Tuple[int] = None,
  bin_size: int = None,
  ):
- if isinstance(self.obs, Projector):
- # branch specifically to handle the projector observable
+ if isinstance(self.obs, BasisStateProjector):
+ # branch specifically to handle the basis state projector observable
  idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
  probs = qml.probs(wires=self.wires).process_samples(
  samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
@@ -122,8 +122,8 @@ def process_samples(
  return qml.math.squeeze(qml.math.mean(samples, axis=axis))
 
  def process_state(self, state: Sequence[complex], wire_order: Wires):
- if isinstance(self.obs, Projector):
- # branch specifically to handle the projector observable
+ if isinstance(self.obs, BasisStateProjector):
+ # branch specifically to handle the basis state projector observable
  idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
  probs = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
  return probs[idx]

pennylane/measurements/var.py

@@ -20,7 +20,7 @@
 
 import pennylane as qml
 from pennylane.operation import Operator
-from pennylane.ops import Projector
+from pennylane.ops.qubit.observables import BasisStateProjector
 from pennylane.wires import Wires
 
 from .measurements import SampleMeasurement, StateMeasurement, Variance
@@ -104,8 +104,8 @@ def process_samples(
  shot_range: Tuple[int] = None,
  bin_size: int = None,
  ):
- if isinstance(self.obs, Projector):
- # branch specifically to handle the projector observable
+ if isinstance(self.obs, BasisStateProjector):
+ # branch specifically to handle the basis state projector observable
  idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
  # we use ``self.wires`` instead of ``self.obs`` because the observable was
  # already applied before the sampling
@@ -125,8 +125,8 @@ def process_samples(
  return qml.math.squeeze(qml.math.var(samples, axis=axis))
 
  def process_state(self, state: Sequence[complex], wire_order: Wires):
- if isinstance(self.obs, Projector):
- # branch specifically to handle the projector observable
+ if isinstance(self.obs, BasisStateProjector):
+ # branch specifically to handle the basis state projector observable
  idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
  # we use ``self.wires`` instead of ``self.obs`` because the observable was
  # already applied to the state

pennylane/ops/qubit/observables.py

@@ -365,11 +365,10 @@ class Projector(Observable):
  0.25
 
  """
+ name = "Projector"
  num_wires = AnyWires
  num_params = 1
  """int: Number of trainable parameters that the operator depends on."""
- _basis_state_type = None # type if Projector inherits from _BasisStateProjector
- _state_vector_type = None # type if Projector inherits from _StateVectorProjector
 
  def __new__(cls, state, wires, **_):
  """Changes parents based on the state representation.
@@ -396,16 +395,10 @@ def __new__(cls, state, wires, **_):
  raise ValueError(f"Input state must be one-dimensional; got shape {shape}.")
 
  if len(state) == len(wires):
- if cls._basis_state_type is None:
- base_cls = (_BasisStateProjector, Projector)
- cls._basis_state_type = type("Projector", base_cls, dict(cls.__dict__))
- return object.__new__(cls._basis_state_type)
+ return object.__new__(BasisStateProjector)
 
  if len(state) == 2 ** len(wires):
- if cls._state_vector_type is None:
- base_cls = (_StateVectorProjector, Projector)
- cls._state_vector_type = type("Projector", base_cls, dict(cls.__dict__))
- return object.__new__(cls._state_vector_type)
+ return object.__new__(StateVectorProjector)
 
  raise ValueError(
  "Input state should have the same length as the wires in the case "
@@ -418,28 +411,25 @@ def pow(self, z):
  """Raise this projector to the power ``z``."""
  return [copy(self)] if (isinstance(z, int) and z > 0) else super().pow(z)
 
- def __copy__(self):
- copied_op = self.__new__(Projector, self.data[0], self.wires)
- copied_op.data = copy(self.data)
- for attr, value in vars(self).items():
- if attr != "data":
- setattr(copied_op, attr, value)
 
- return copied_op
+class BasisStateProjector(Projector):
+ r"""Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`, where
+ :math:`\phi` denotes a basis state."""
 
-
-class _BasisStateProjector(Observable):
  # The call signature should be the same as Projector.__new__ for the positional
  # arguments, but with free key word arguments.
  def __init__(self, state, wires, id=None):
  wires = Wires(wires)
- state = list(qml.math.toarray(state))
+ state = list(qml.math.toarray(state).astype(int))
 
  if not set(state).issubset({0, 1}):
  raise ValueError(f"Basis state must only consist of 0s and 1s; got {state}")
 
  super().__init__(state, wires=wires, id=id)
 
+ def __new__(cls): # pylint: disable=arguments-differ
+ return object.__new__(cls)
+
  def label(self, decimals=None, base_label=None, cache=None):
  r"""A customizable string representation of the operator.
 
@@ -455,7 +445,7 @@ def label(self, decimals=None, base_label=None, cache=None):
 
  **Example:**
 
- >>> _BasisStateProjector([0, 1, 0], wires=(0, 1, 2)).label()
+ >>> BasisStateProjector([0, 1, 0], wires=(0, 1, 2)).label()
  '|010⟩⟨010|'
 
  """
@@ -472,7 +462,7 @@ def compute_matrix(basis_state): # pylint: disable=arguments-differ
  The canonical matrix is the textbook matrix representation that does not consider wires.
  Implicitly, this assumes that the wires of the operator correspond to the global wire order.
 
- .. seealso:: :meth:`~._BasisStateProjector.matrix`
+ .. seealso:: :meth:`~.BasisStateProjector.matrix`
 
  Args:
  basis_state (Iterable): basis state to project on
@@ -482,7 +472,7 @@ def compute_matrix(basis_state): # pylint: disable=arguments-differ
 
  **Example**
 
- >>> _BasisStateProjector.compute_matrix([0, 1])
+ >>> BasisStateProjector.compute_matrix([0, 1])
  [[0. 0. 0. 0.]
  [0. 1. 0. 0.]
  [0. 0. 0. 0.]
@@ -506,7 +496,7 @@ def compute_eigvals(basis_state): # pylint: disable=arguments-differ
 
  Otherwise, no particular order for the eigenvalues is guaranteed.
 
- .. seealso:: :meth:`~._BasisStateProjector.eigvals`
+ .. seealso:: :meth:`~.BasisStateProjector.eigvals`
 
  Args:
  basis_state (Iterable): basis state to project on
@@ -516,7 +506,7 @@ def compute_eigvals(basis_state): # pylint: disable=arguments-differ
 
  **Example**
 
- >>> _BasisStateProjector.compute_eigvals([0, 1])
+ >>> BasisStateProjector.compute_eigvals([0, 1])
  [0. 1. 0. 0.]
  """
  w = np.zeros(2 ** len(basis_state))
@@ -537,7 +527,7 @@ def compute_diagonalizing_gates(
  The diagonalizing gates rotate the state into the eigenbasis
  of the operator.
 
- .. seealso:: :meth:`~._BasisStateProjector.diagonalizing_gates`.
+ .. seealso:: :meth:`~.BasisStateProjector.diagonalizing_gates`.
 
  Args:
  basis_state (Iterable): basis state that the operator projects on
@@ -547,13 +537,16 @@ def compute_diagonalizing_gates(
 
  **Example**
 
- >>> _BasisStateProjector.compute_diagonalizing_gates([0, 1, 0, 0], wires=[0, 1])
+ >>> BasisStateProjector.compute_diagonalizing_gates([0, 1, 0, 0], wires=[0, 1])
  []
  """
  return []
 
 
-class _StateVectorProjector(Observable):
+class StateVectorProjector(Projector):
+ r"""Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`, where
+ :math:`\phi` denotes a state."""
+
  # The call signature should be the same as Projector.__new__ for the positional
  # arguments, but with free key word arguments.
  def __init__(self, state, wires, id=None):
@@ -562,6 +555,9 @@ def __init__(self, state, wires, id=None):
 
  super().__init__(state, wires=wires, id=id)
 
+ def __new__(cls): # pylint: disable=arguments-differ
+ return object.__new__(cls)
+
  def label(self, decimals=None, base_label=None, cache=None):
  r"""A customizable string representation of the operator.
 
@@ -578,14 +574,14 @@ def label(self, decimals=None, base_label=None, cache=None):
  **Example:**
 
  >>> state_vector = np.array([0, 1, 1, 0])/np.sqrt(2)
- >>> _StateVectorProjector(state_vector, wires=(0, 1)).label()
+ >>> StateVectorProjector(state_vector, wires=(0, 1)).label()
  'P'
- >>> _StateVectorProjector(state_vector, wires=(0, 1)).label(base_label="hi!")
+ >>> StateVectorProjector(state_vector, wires=(0, 1)).label(base_label="hi!")
  'hi!'
  >>> dev = qml.device("default.qubit", wires=1)
  >>> @qml.qnode(dev)
  >>> def circuit(state):
- ... return qml.expval(_StateVectorProjector(state, [0]))
+ ... return qml.expval(StateVectorProjector(state, [0]))
  >>> print(qml.draw(circuit)([1, 0]))
  0: ───┤ <|0⟩⟨0|>
  >>> print(qml.draw(circuit)(np.array([1, 1]) / np.sqrt(2)))
@@ -631,7 +627,7 @@ def compute_matrix(state_vector): # pylint: disable=arguments-differ,arguments-
 
  The projector of the state :math:`\frac{1}{\sqrt{2}}(\ket{01}+\ket{10})`
 
- >>> _StateVectorProjector.compute_matrix([0, 1/np.sqrt(2), 1/np.sqrt(2), 0])
+ >>> StateVectorProjector.compute_matrix([0, 1/np.sqrt(2), 1/np.sqrt(2), 0])
  [[0. 0. 0. 0.]
  [0. 0.5 0.5 0.]
  [0. 0.5 0.5 0.]
@@ -652,7 +648,7 @@ def compute_eigvals(state_vector): # pylint: disable=arguments-differ,arguments
 
  Otherwise, no particular order for the eigenvalues is guaranteed.
 
- .. seealso:: :meth:`~._StateVectorProjector.eigvals`
+ .. seealso:: :meth:`~.StateVectorProjector.eigvals`
 
  Args:
  state_vector (Iterable): state vector to project on
@@ -662,7 +658,7 @@ def compute_eigvals(state_vector): # pylint: disable=arguments-differ,arguments
 
  **Example**
 
- >>> _StateVectorProjector.compute_eigvals([0, 0, 1, 0])
+ >>> StateVectorProjector.compute_eigvals([0, 0, 1, 0])
  array([1, 0, 0, 0])
  """
  w = qml.math.zeros_like(state_vector)
@@ -682,7 +678,7 @@ def compute_diagonalizing_gates(
  The diagonalizing gates rotate the state into the eigenbasis
  of the operator.
 
- .. seealso:: :meth:`~._StateVectorProjector.diagonalizing_gates`.
+ .. seealso:: :meth:`~.StateVectorProjector.diagonalizing_gates`.
 
  Args:
  state_vector (Iterable): state vector that the operator projects on.
@@ -693,7 +689,7 @@ def compute_diagonalizing_gates(
  **Example**
 
  >>> state_vector = np.array([1., 1j])/np.sqrt(2)
- >>> _StateVectorProjector.compute_diagonalizing_gates(state_vector, wires=[0])
+ >>> StateVectorProjector.compute_diagonalizing_gates(state_vector, wires=[0])
  [QubitUnitary(array([[ 0.70710678+0.j , 0. -0.70710678j],
  [ 0. +0.70710678j, -0.70710678+0.j ]]), wires=[0])]
  """

tests/devices/experimental/test_default_qubit_2.py

@@ -1787,6 +1787,26 @@ def test_shot_vectors(self, max_workers, n_qubits, shots):
  assert np.all(np.logical_or(np.logical_or(r[1] == 0, r[1] == 1), r[1] == 2))
 
 
+class TestDynamicType:
+ """Tests the compatibility with dynamic type classes such as `qml.Projector`."""
+
+ @pytest.mark.parametrize("n_wires", [1, 2, 3])
+ @pytest.mark.parametrize("max_workers", [None, 1, 2])
+ def test_projector(self, max_workers, n_wires):
+ """Test that qml.Projector yields the expected results for both of its subclasses."""
+ wires = list(range(n_wires))
+ dev = DefaultQubit2(max_workers=max_workers)
+ ops = [qml.Hadamard(q) for q in wires]
+ basis_state = np.zeros((n_wires,))
+ state_vector = np.zeros((2**n_wires,))
+ state_vector[0] = 1
+
+ for state in [basis_state, state_vector]:
+ qs = qml.tape.QuantumScript(ops, [qml.expval(qml.Projector(state, wires))])
+ res = dev.execute(qs)
+ assert np.isclose(res, 1 / 2**n_wires)
+
+
 @pytest.mark.parametrize("max_workers", [None, 1, 2])
 def test_broadcasted_parameter(max_workers):
  """Test that DefaultQubit2 handles broadcasted parameters as expected."""

tests/gradients/parameter_shift/test_parameter_shift.py

@@ -2183,17 +2183,17 @@ def test_recycling_unshifted_tape_result(self):
  # + 2 operations x 2 shifted positions + 1 unshifted term <-- <H^2>
  assert len(tapes) == (2 * 2 + 1) + (2 * 2 + 1)
 
- def test_projector_variance(self, tol):
+ @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector
+ def test_projector_variance(self, state, tol):
  """Test that the variance of a projector is correctly returned"""
  dev = qml.device("default.qubit", wires=2)
- P = np.array([1])
  x, y = 0.765, -0.654
 
  with qml.queuing.AnnotatedQueue() as q:
  qml.RX(x, wires=0)
  qml.RY(y, wires=1)
  qml.CNOT(wires=[0, 1])
- qml.var(qml.Projector(P, wires=0) @ qml.PauliX(1))
+ qml.var(qml.Projector(state, wires=0) @ qml.PauliX(1))
 
  tape = qml.tape.QuantumScript.from_queue(q)
  tape.trainable_params = {0, 1}
@@ -2983,17 +2983,17 @@ def test_expval_and_variance(self, tol):
  # assert gradA == pytest.approx(expected, abs=tol)
  # assert gradF == pytest.approx(expected, abs=tol)
 
- def test_projector_variance(self, tol):
+ @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector
+ def test_projector_variance(self, state, tol):
  """Test that the variance of a projector is correctly returned"""
  dev = qml.device("default.qubit", wires=2)
- P = np.array([1])
  x, y = 0.765, -0.654
 
  with qml.queuing.AnnotatedQueue() as q:
  qml.RX(x, wires=0)
  qml.RY(y, wires=1)
  qml.CNOT(wires=[0, 1])
- qml.var(qml.Projector(P, wires=0) @ qml.PauliX(1))
+ qml.var(qml.Projector(state, wires=0) @ qml.PauliX(1))
 
  tape = qml.tape.QuantumScript.from_queue(q)
  tape.trainable_params = {0, 1}