Upgrade and generalise basis state preparation

doc/development/deprecations.rst

@@ -85,6 +85,12 @@ Pending deprecations
   - Deprecated in v0.37
   - Will be removed in v0.39
 
+* The ``BasisStatePreparation`` template is deprecated.
+  Instead, use ``BasisState``.
+
+  - Deprecated in v0.38
+  - Will be removed in v0.39
+
 New operator arithmetic deprecations
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

doc/releases/changelog-dev.md

@@ -260,6 +260,10 @@
   Instead, use `pennylane.gradients.classical_fisher` and `pennylane.gradients.quantum_fisher`.
   [(#5985)](https://github.com/PennyLaneAI/pennylane/pull/5985)
 
+* The ``BasisStatePreparation`` template is deprecated.
+  Instead, ``BasisState`` can be called on the constructed operator.
+  [(#6021)](https://github.com/PennyLaneAI/pennylane/pull/6021)
+
 * The legacy devices `default.qubit.{autograd,torch,tf,jax,legacy}` are deprecated.
   Instead, use `default.qubit` as it now supports backpropagation through the several backends.
   [(#5997)](https://github.com/PennyLaneAI/pennylane/pull/5997)
@@ -310,7 +314,16 @@
   [(#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.BasisState` now works with jax.jit.
+  [(#6021)](https://github.com/PennyLaneAI/pennylane/pull/6021)
+
+* `qml.BasisEmbedding` now gives the correct decomposition.
+  [(#6021)](https://github.com/PennyLaneAI/pennylane/pull/6021)
+
+* `qml.BasisEmbedding` now works with lightning.qubit and jax.jit.
+  [(#6021)](https://github.com/PennyLaneAI/pennylane/pull/6021)
 
 * Jacobian shape is fixed for measurements with dimension in `qml.gradients.vjp.compute_vjp_single`.
 [(5986)](https://github.com/PennyLaneAI/pennylane/pull/5986)

pennylane/devices/tests/test_templates.py

@@ -225,6 +225,15 @@ def circuit():
             [math.fidelity_statevector(circuit(), exp_state)], [1.0], atol=tol(dev.shots)
         )
 
+    @pytest.fixture(scope="function", autouse=True)
+    def capture_warnings(self, recwarn):
+        """Capture warnings."""
+        yield
+        if len(recwarn) > 0:
+            for w in recwarn:
+                assert isinstance(w.message, qml.PennyLaneDeprecationWarning)
+                assert "BasisStatePreparation is deprecated" in str(w.message)
+
     def test_BasisStatePreparation(self, device, tol):
         """Test the BasisStatePreparation template."""
         dev = device(4)

pennylane/ops/qubit/state_preparation.py

@@ -18,23 +18,23 @@
 # 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"}
 
 
 class BasisState(StatePrepBase):
-    r"""BasisState(n, wires)
+    r"""BasisState(features, wires)
     Prepares a single computational basis state.
 
     **Details:**
 
     * Number of wires: Any (the operation can act on any number of wires)
     * Number of parameters: 1
-    * Gradient recipe: None (integer parameters not supported)
 
     .. note::
 
@@ -48,9 +48,8 @@ class BasisState(StatePrepBase):
         as :math:`U|0\rangle = |\psi\rangle`
 
     Args:
-        n (array): prepares the basis state :math:`\ket{n}`, where ``n`` is an
-            array of integers from the set :math:`\{0, 1\}`, i.e.,
-            if ``n = np.array([0, 1, 0])``, prepares the state :math:`|010\rangle`.
+        features (tensor_like): binary input of shape ``(len(wires), )``. For example, for ``features=np.array([0, 1, 0])`` or ``features=2`` (binary 010), the quantum system will be prepared in state :math:`|010 \rangle`.
+
         wires (Sequence[int] or int): the wire(s) the operation acts on
         id (str): custom label given to an operator instance,
             can be useful for some applications where the instance has to be identified.
@@ -66,15 +65,51 @@ 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):
 
-    ndim_params = (1,)
-    """int: Number of dimensions per trainable parameter of the operator."""
+        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 ** 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}")
+
+        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.
@@ -93,33 +128,49 @@ def compute_decomposition(n, wires):
         **Example:**
 
         >>> qml.BasisState.compute_decomposition([1,0], wires=(0,1))
-        [BasisStatePreparation([1, 0], wires=[0, 1])]
+        [X(0)]
 
         """
-        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 (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/templates/embeddings/basis.py

@@ -15,17 +15,15 @@
 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):
+# pylint: disable=missing-class-docstring
+class BasisEmbedding(BasisState):
     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
+    For example, for ``features=np.array([0, 1, 0])`` or ``features=2`` (binary 010), the
     quantum system will be prepared in state :math:`|010 \rangle`.
 
     .. warning::
@@ -35,8 +33,9 @@ class BasisEmbedding(Operation):
         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
+        features (tensor_like or int): binary input of shape ``(len(wires), )`` or integer
+            that represents the binary input.
+        wires (Any or Iterable[Any]): wires that the template acts on.
 
     Example:
 
@@ -68,93 +67,3 @@ def circuit(feature_vector):
         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

pennylane/templates/state_preparations/basis.py

@@ -15,6 +15,8 @@
 Contains the BasisStatePreparation template.
 """
 
+import warnings
+
 import numpy as np
 
 import pennylane as qml
@@ -30,6 +32,8 @@ class BasisStatePreparation(Operation):
         ``basis_state`` influences the circuit architecture and is therefore incompatible with
         gradient computations.
 
+        ``BasisStatePreparation`` is deprecated and will be removed in version 0.39. Instead, please use ``BasisState``.
+
     Args:
         basis_state (array): Input array of shape ``(n,)``, where n is the number of wires
             the state preparation acts on.
@@ -59,6 +63,13 @@ def circuit(basis_state):
     ndim_params = (1,)
 
     def __init__(self, basis_state, wires, id=None):
+
+        warnings.warn(
+            "BasisStatePreparation is deprecated and will be removed in version 0.39. "
+            "Instead, please use BasisState.",
+            qml.PennyLaneDeprecationWarning,
+        )
+
         basis_state = qml.math.stack(basis_state)
 
         # check if the `basis_state` param is batched

tests/capture/test_templates.py

@@ -190,6 +190,15 @@ def enable_disable_plxpr():
 ]
 
 
+@pytest.fixture(scope="function", autouse=True)
+def capture_warnings(recwarn):
+    yield
+    if len(recwarn) > 0:
+        for w in recwarn:
+            assert isinstance(w.message, qml.PennyLaneDeprecationWarning)
+            assert "BasisStatePreparation is deprecated" in str(w.message)
+
+
 @pytest.mark.parametrize("template, args, kwargs", unmodified_templates_cases)
 def test_unmodified_templates(template, args, kwargs):
     """Test that templates with unmodified primitive binds are captured as expected."""