Upgrade and generalise basis state preparation

pennylane/ops/qubit/state_preparation.py

@@ -18,9 +18,10 @@
 # pylint:disable=abstract-method,arguments-differ,protected-access,no-member
 import numpy as np
 
+import pennylane as qml
 from pennylane import math
 from pennylane.operation import AnyWires, Operation, StatePrepBase
-from pennylane.templates.state_preparations import BasisStatePreparation, MottonenStatePreparation
+from pennylane.templates.state_preparations import MottonenStatePreparation
 from pennylane.wires import WireError, Wires
 
 state_prep_ops = {"BasisState", "StatePrep", "QubitDensityMatrix"}
@@ -66,15 +67,57 @@ class BasisState(StatePrepBase):
     [0.+0.j 0.+0.j 0.+0.j 1.+0.j]
     """
 
-    num_wires = AnyWires
-    num_params = 1
-    """int: Number of trainable parameters that the operator depends on."""
+    def __init__(self, features, wires, id=None):
+        if isinstance(features, list):
+            features = qml.math.stack(features)
 
-    ndim_params = (1,)
-    """int: Number of dimensions per trainable parameter of the operator."""
+        tracing = qml.math.is_abstract(features)
+
+        if qml.math.shape(features) == ():
+            if not tracing and features >= 2 ** len(wires):
+                raise ValueError(
+                    f"Features must be of length {len(wires)}, got features={features} which is >= {2 ** len(wires)}"
+                )
+            bin = 2 ** math.arange(len(wires))[::-1]
+            features = qml.math.where((features & bin) > 0, 1, 0)
+
+        wires = Wires(wires)
+        shape = qml.math.shape(features)
+
+        if len(shape) != 1:
+            raise ValueError(f"Features must be one-dimensional; got shape {shape}.")
+
+        n_features = shape[0]
+        if n_features != len(wires):
+            raise ValueError(
+                f"Features must be of length {len(wires)}; got length {n_features} (features={features})."
+            )
+
+        if not tracing:
+            features_list = list(qml.math.toarray(features))
+            if not set(features_list).issubset({0, 1}):
+                raise ValueError(f"Basis state must only consist of 0s and 1s; got {features_list}")
+
+        if len(qml.math.shape(features)) > 1:
+            raise ValueError(
+                "Broadcasting with BasisState is not supported. Please use the "
+                "qml.transforms.broadcast_expand transform to use broadcasting with "
+                "BasisState."
+            )
+
+        super().__init__(features, wires=wires, id=id)
+
+    def _flatten(self):
+        features = self.parameters[0]
+        features = tuple(features) if isinstance(features, list) else features
+        return (features,), (self.wires,)
+
+    @classmethod
+    def _unflatten(cls, data, metadata) -> "BasisState":
+        return cls(data[0], wires=metadata[0])
 
     @staticmethod
-    def compute_decomposition(n, wires):
+    def compute_decomposition(features, wires):
         r"""Representation of the operator as a product of other operators (static method). :
 
         .. math:: O = O_1 O_2 \dots O_n.
@@ -96,30 +139,50 @@ def compute_decomposition(n, wires):
         [BasisStatePreparation([1, 0], wires=[0, 1])]
 
         """
-        return [BasisStatePreparation(n, wires)]
+
+        if not qml.math.is_abstract(features):
+            op_list = []
+            for wire, state in zip(wires, features):
+                if state == 1:
+                    op_list.append(qml.X(wire))
+            return op_list
+
+        op_list = []
+        for wire, state in zip(wires, features):
+            op_list.append(qml.PhaseShift(state * np.pi / 2, wire))
+            op_list.append(qml.RX(state * np.pi, wire))
+            op_list.append(qml.PhaseShift(state * np.pi / 2, wire))
+
+        return op_list
 
     def state_vector(self, wire_order=None):
         """Returns a statevector of shape ``(2,) * num_wires``."""
         prep_vals = self.parameters[0]
-        if any(i not in [0, 1] for i in prep_vals):
-            raise ValueError("BasisState parameter must consist of 0 or 1 integers.")
+        if qml.math.shape(prep_vals) == ():
+            bin = 2 ** math.arange(len(self.wires))[::-1]
+            prep_vals = qml.math.where((prep_vals & bin) > 0, 1, 0)
 
-        if (num_wires := len(self.wires)) != len(prep_vals):
-            raise ValueError("BasisState parameter and wires must be of equal length.")
+        prep_vals_int = math.cast(prep_vals, int)
 
-        prep_vals = math.cast(prep_vals, int)
         if wire_order is None:
-            indices = prep_vals
+            indices = prep_vals_int
+            num_wires = len(indices)
         else:
             if not Wires(wire_order).contains_wires(self.wires):
                 raise WireError("Custom wire_order must contain all BasisState wires")
             num_wires = len(wire_order)
             indices = [0] * num_wires
-            for base_wire_label, value in zip(self.wires, prep_vals):
+            for base_wire_label, value in zip(self.wires, prep_vals_int):
                 indices[wire_order.index(base_wire_label)] = value
 
-        ket = np.zeros((2,) * num_wires)
-        ket[tuple(indices)] = 1
+        if qml.math.get_interface(prep_vals_int) == "jax":
+            ket = math.array(math.zeros((2,) * num_wires), like="jax")
+            ket = ket.at[tuple(indices)].set(1)
+
+        else:
+            ket = math.zeros((2,) * num_wires)
+            ket[tuple(indices)] = 1
+
         return math.convert_like(ket, prep_vals)
 
 

pennylane/optimize/spsa.py

@@ -270,7 +270,7 @@ def compute_grad(self, objective_fn, args, kwargs):
                 shots = Shots(objective_fn.device._raw_shot_sequence)  # pragma: no cover
             else:
                 shots = Shots(None)
-            if np.prod(objective_fn.func(*args).shape(objective_fn.device, shots)) > 1:
+            if np.prod(objective_fn.func(*args, **kwargs).shape(objective_fn.device, shots)) > 1:
                 raise ValueError(
                     "The objective function must be a scalar function for the gradient "
                     "to be computed."

pennylane/templates/embeddings/basis.py

@@ -15,146 +15,10 @@
 Contains the BasisEmbedding template.
 """
 # pylint: disable-msg=too-many-branches,too-many-arguments,protected-access
-import numpy as np
 
-import pennylane as qml
-from pennylane.operation import AnyWires, Operation
-from pennylane.wires import Wires
+from pennylane.ops.qubit.state_preparation import BasisState
 
 
-class BasisEmbedding(Operation):
-    r"""Encodes :math:`n` binary features into a basis state of :math:`n` qubits.
-
-    For example, for ``features=np.array([0, 1, 0])`` or ``features=2`` (binary 10), the
-    quantum system will be prepared in state :math:`|010 \rangle`.
-
-    .. warning::
-
-        ``BasisEmbedding`` calls a circuit whose architecture depends on the binary features.
-        The ``features`` argument is therefore not differentiable when using the template, and
-        gradients with respect to the argument cannot be computed by PennyLane.
-
-    Args:
-        features (tensor_like): binary input of shape ``(len(wires), )``
-        wires (Any or Iterable[Any]): wires that the template acts on
-
-    Example:
-
-        Basis embedding encodes the binary feature vector into a basis state.
-
-        .. code-block:: python
-
-            dev = qml.device('default.qubit', wires=3)
-
-            @qml.qnode(dev)
-            def circuit(feature_vector):
-                qml.BasisEmbedding(features=feature_vector, wires=range(3))
-                return qml.state()
-
-            X = [1,1,1]
-
-        The resulting circuit is:
-
-        >>> print(qml.draw(circuit, level="device")(X))
-        0: ──X─┤  State
-        1: ──X─┤  State
-        2: ──X─┤  State
-
-        And, the output state is:
-
-        >>> print(circuit(X))
-            [0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 1.+0.j]
-
-        Thus, ``[1,1,1]`` is mapped to :math:`|111 \rangle`.
-
-    """
-
-    num_wires = AnyWires
-    grad_method = None
-
-    def _flatten(self):
-        basis_state = self.hyperparameters["basis_state"]
-        basis_state = tuple(basis_state) if isinstance(basis_state, list) else basis_state
-        return tuple(), (self.wires, basis_state)
-
-    @classmethod
-    def _unflatten(cls, _, metadata) -> "BasisEmbedding":
-        return cls(features=metadata[1], wires=metadata[0])
-
-    def __init__(self, features, wires, id=None):
-        if isinstance(features, list):
-            features = qml.math.stack(features)
-
-        tracing = qml.math.is_abstract(features)
-
-        if qml.math.shape(features) == ():
-            if not tracing and features >= 2 ** len(wires):
-                raise ValueError(
-                    f"Features must be of length {len(wires)}, got features={features} which is >= {2 ** len(wires)}"
-                )
-            bin = 2 ** np.arange(len(wires))[::-1]
-            features = qml.math.where((features & bin) > 0, 1, 0)
-
-        wires = Wires(wires)
-        shape = qml.math.shape(features)
-
-        if len(shape) != 1:
-            raise ValueError(f"Features must be one-dimensional; got shape {shape}.")
-
-        n_features = shape[0]
-        if n_features != len(wires):
-            raise ValueError(
-                f"Features must be of length {len(wires)}; got length {n_features} (features={features})."
-            )
-
-        if not tracing:
-            features = list(qml.math.toarray(features))
-            if not set(features).issubset({0, 1}):
-                raise ValueError(f"Basis state must only consist of 0s and 1s; got {features}")
-
-        self._hyperparameters = {"basis_state": features}
-
-        super().__init__(wires=wires, id=id)
-
-    @property
-    def num_params(self):
-        return 0
-
-    @staticmethod
-    def compute_decomposition(wires, basis_state):  # pylint: disable=arguments-differ
-        r"""Representation of the operator as a product of other operators.
-
-        .. math:: O = O_1 O_2 \dots O_n.
-
-
-
-        .. seealso:: :meth:`~.BasisEmbedding.decomposition`.
-
-        Args:
-            features (tensor-like): binary input of shape ``(len(wires), )``
-            wires (Any or Iterable[Any]): wires that the operator acts on
-
-        Returns:
-            list[.Operator]: decomposition of the operator
-
-        **Example**
-
-        >>> features = torch.tensor([1, 0, 1])
-        >>> qml.BasisEmbedding.compute_decomposition(features, wires=["a", "b", "c"])
-        [X('a'),
-         X('c')]
-        """
-        if not qml.math.is_abstract(basis_state):
-            ops_list = []
-            for wire, bit in zip(wires, basis_state):
-                if bit == 1:
-                    ops_list.append(qml.X(wire))
-            return ops_list
-
-        ops_list = []
-        for wire, state in zip(wires, basis_state):
-            ops_list.append(qml.PhaseShift(state * np.pi / 2, wire))
-            ops_list.append(qml.RX(state * np.pi, wire))
-            ops_list.append(qml.PhaseShift(state * np.pi / 2, wire))
-
-        return ops_list
+# pylint: disable=missing-class-docstring
+class BasisEmbedding(BasisState):
+    pass  # BasisEmbedding is still available

tests/devices/test_default_qubit_legacy.py

@@ -650,13 +650,14 @@ def test_apply_errors_qubit_state_vector(self, qubit_device_2_wires):
             )
 
     def test_apply_errors_basis_state(self, qubit_device_2_wires):
+
         with pytest.raises(
-            ValueError, match="BasisState parameter must consist of 0 or 1 integers."
+            ValueError, match=r"Basis state must only consist of 0s and 1s; got \[-0\.2, 4\.2\]"
         ):
             qubit_device_2_wires.apply([qml.BasisState(np.array([-0.2, 4.2]), wires=[0, 1])])
 
         with pytest.raises(
-            ValueError, match="BasisState parameter and wires must be of equal length."
+            ValueError, match=r"Features must be of length 1; got length 2 \(features=\[0 1\]\)\."
         ):
             qubit_device_2_wires.apply([qml.BasisState(np.array([0, 1]), wires=[0])])
 

tests/devices/test_default_qubit_tf.py

@@ -295,7 +295,8 @@ def test_invalid_basis_state_length(self):
         state = np.array([0, 0, 1, 0])
 
         with pytest.raises(
-            ValueError, match=r"BasisState parameter and wires must be of equal length"
+            ValueError,
+            match=r"Features must be of length 3; got length 4 \(features=\[0 0 1 0\]\)",
         ):
             dev.apply([qml.BasisState(state, wires=[0, 1, 2])])
 
@@ -305,7 +306,7 @@ def test_invalid_basis_state(self):
         state = np.array([0, 0, 1, 2])
 
         with pytest.raises(
-            ValueError, match=r"BasisState parameter must consist of 0 or 1 integers"
+            ValueError, match=r"Basis state must only consist of 0s and 1s; got \[0, 0, 1, 2\]"
         ):
             dev.apply([qml.BasisState(state, wires=[0, 1, 2, 3])])
 

tests/devices/test_default_qubit_torch.py

@@ -258,7 +258,8 @@ def test_invalid_basis_state_length(self, device, torch_device):
         state = torch.tensor([0, 0, 1, 0])
 
         with pytest.raises(
-            ValueError, match=r"BasisState parameter and wires must be of equal length"
+            ValueError,
+            match=r"Features must be of length 3; got length 4 \(features=tensor\(\[0, 0, 1, 0\]\)\)",
         ):
             dev.apply([qml.BasisState(state, wires=[0, 1, 2])])
 
@@ -268,7 +269,7 @@ def test_invalid_basis_state(self, device, torch_device):
         state = torch.tensor([0, 0, 1, 2])
 
         with pytest.raises(
-            ValueError, match=r"BasisState parameter must consist of 0 or 1 integers"
+            ValueError, match=r"Basis state must only consist of 0s and 1s; got \[0, 0, 1, 2\]"
         ):
             dev.apply([qml.BasisState(state, wires=[0, 1, 2, 3])])
 

tests/ops/qubit/test_state_prep.py

@@ -26,7 +26,7 @@
 @pytest.mark.parametrize(
     "op",
     [
-        qml.BasisState(np.array([0, 1]), wires=0),
+        qml.BasisState(np.array([0, 1]), wires=[0, 1]),
         qml.StatePrep(np.array([1.0, 0.0]), wires=0),
         qml.QubitDensityMatrix(densitymat0, wires=0),
     ],
@@ -66,8 +66,6 @@ def test_BasisState_decomposition(self):
         ops2 = qml.BasisState(n, wires=wires).decomposition()
 
         assert len(ops1) == len(ops2) == 1
-        assert isinstance(ops1[0], qml.BasisStatePreparation)
-        assert isinstance(ops2[0], qml.BasisStatePreparation)
 
     def test_StatePrep_decomposition(self):
         """Test the decomposition for StatePrep."""
@@ -351,18 +349,9 @@ def test_BasisState_state_vector_bad_wire_order(self):
         with pytest.raises(WireError, match="wire_order must contain all BasisState wires"):
             basis_op.state_vector(wire_order=[1, 2])
 
-    def test_BasisState_explicitly_checks_0_1(self):
-        """Tests that BasisState gives a clear error if a value other than 0 or 1 is given."""
-        op = qml.BasisState([2, 1], wires=[0, 1])
-        with pytest.raises(
-            ValueError, match="BasisState parameter must consist of 0 or 1 integers."
-        ):
-            _ = op.state_vector()
-
     def test_BasisState_wrong_param_size(self):
         """Tests that the parameter must be of length num_wires."""
-        op = qml.BasisState([0], wires=[0, 1])
         with pytest.raises(
-            ValueError, match="BasisState parameter and wires must be of equal length."
+            ValueError, match=r"Features must be of length 2; got length 1 \(features=\[0\]\)."
         ):
-            _ = op.state_vector()
+            _ = qml.BasisState([0], wires=[0, 1])

tests/tape/test_qscript.py

@@ -641,7 +641,7 @@ def test_deep_copy(self):
 def test_adjoint():
     """Tests taking the adjoint of a quantum script."""
     ops = [
-        qml.BasisState([1, 1], wires=0),
+        qml.BasisState([1, 1], wires=[0, 1]),
         qml.RX(1.2, wires=0),
         qml.S(0),
         qml.CNOT((0, 1)),