Fix stateprep casting rules. (#501)

* Cast stateprep to complex128 in doubt.

* Auto update version

* Auto update version

* trigger CI

* Special case integers.

* Fix test_comparison.

* Update changelog.

* One liner return.

* Update .github/CHANGELOG.md

Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>

* Try Christina's solution.

* Revert to _preprocess_state_vector fix and add unit test.

* Auto update version

* trigger CI

---------

Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>

.github/CHANGELOG.md

@@ -7,15 +7,18 @@
 
 ### Breaking changes
 
+* Cast integral-valued arrays to the device's complex type on entry in `_preprocess_state_vector` to ensure the state is correctly represented with floating-point numbers.  
+  [(#501)](https://github.com/PennyLaneAI/pennylane-lightning/pull/501)
+
+* Update DefaultQubit to DefaultQubitLegacy on Lightning fallback.
+  [(#500)](https://github.com/PennyLaneAI/pennylane-lightning/pull/500)
+
 * Enums defined in `GateOperation.hpp` start at `1` (previously `0`). `::BEGIN` is introduced in a few places where it was assumed `0` accordingly.
   [(#485)](https://github.com/PennyLaneAI/pennylane-lightning/pull/485)
 
 * Enable pre-commit hooks to format all Python files and linting of all Python source files.
   [(#485)](https://github.com/PennyLaneAI/pennylane-lightning/pull/485)
 
-* Update DefaultQubit to DefaultQubitLegacy on Lightning fallback.
-  [(#500)](https://github.com/PennyLaneAI/pennylane-lightning/pull/500)
-
 ### Improvements
 
 * Refactor LKokkos `StateVectorKokkos` class to use Kokkos `RangePolicy` together with special functors in `applyMultiQubitOp` to apply 1- to 4-wire generic unitary gates. For more than 4 wires, the general implementation using Kokkos `TeamPolicy` is employed to yield the best all-around performance.

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.33.0-dev10"
+__version__ = "0.33.0-dev11"

pennylane_lightning/core/lightning_base.py

@@ -103,6 +103,7 @@ def accepts_obj(obj):
 
         return qml.BooleanFn(accepts_obj)
 
+    # pylint: disable=missing-function-docstring
     @classmethod
     def capabilities(cls):
         capabilities = super().capabilities().copy()
@@ -206,6 +207,9 @@ def _preprocess_state_vector(self, state, device_wires):
         # translate to wire labels used by device
         device_wires = self.map_wires(device_wires)
 
+        # special case for integral types
+        if state.dtype.kind == "i":
+            state = qml.numpy.array(state, dtype=self.C_DTYPE)
         state = self._asarray(state, dtype=self.C_DTYPE)
 
         if len(device_wires) == self.num_wires and Wires(sorted(device_wires)) == device_wires:

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -205,6 +205,7 @@ def __init__(
             self._state = self._create_basis_state(0)
             self._pre_rotated_state = self._state
 
+            self._batch_obs = batch_obs
             self._mcmc = mcmc
             if self._mcmc:
                 if kernel_name not in [
@@ -281,7 +282,7 @@ def measurements(self):
 
         @property
         def state(self):
-            # Flattening the state.
+            """Returns the flattened state vector."""
             shape = (1 << self.num_wires,)
             return self._reshape(self._pre_rotated_state, shape)
 
@@ -368,7 +369,9 @@ def apply_lightning(self, state, operations):
 
             return np.reshape(state_vector, state.shape)
 
+        # pylint: disable=unused-argument
         def apply(self, operations, rotations=None, **kwargs):
+            """Applies operations to the state vector."""
             # State preparation is currently done in Python
             if operations:  # make sure operations[0] exists
                 if isinstance(operations[0], StatePrep):
@@ -629,6 +632,7 @@ def _init_process_jacobian_tape(self, tape, starting_state, use_device_state):
             return StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket)
 
         def adjoint_jacobian(self, tape, starting_state=None, use_device_state=False):
+            """Computes and returns the Jacobian with the adjoint method."""
             if self.shots is not None:
                 warn(
                     "Requested adjoint differentiation to be computed with finite shots. "
@@ -812,7 +816,7 @@ def processing_fn_state(tape):
 else:
 
     class LightningQubit(LightningBaseFallBack):  # pragma: no cover
-        # pylint: disable=missing-class-docstring
+        # pylint: disable=missing-class-docstring, too-few-public-methods
         name = "Lightning qubit PennyLane plugin [No binaries found - Fallback: default.qubit]"
         short_name = "lightning.qubit"
 

tests/test_apply.py

@@ -211,6 +211,24 @@ def test_apply_operation_state_preparation(
 
         assert np.allclose(dev.state, np.array(expected_output), atol=tol, rtol=0)
 
+    def test_integer_state_preparation(self, qubit_device, tol):
+        """Tests that applying an operation yields the expected output state for single wire
+        operations that have no parameters."""
+        dev = qubit_device(wires=2)
+
+        @qml.qnode(dev)
+        def circuit0():
+            qml.RX(0.2, wires=[0])
+            return qml.state()
+
+        @qml.qnode(dev)
+        def circuit1():
+            qml.StatePrep(np.array([1, 0, 0, 0]), wires=[0, 1])
+            qml.RX(0.2, wires=[0])
+            return qml.state()
+
+        assert np.allclose(circuit0(), circuit1(), atol=tol, rtol=0)
+
     """ operation,input,expected_output,par """
     test_data_single_wire_with_parameters = [
         (qml.PhaseShift, [1, 0], [1, 0], [math.pi / 2]),

tests/test_comparison.py

@@ -73,24 +73,23 @@ def test_one_qubit_circuit(
 
         monkeypatch.setenv("OMP_NUM_THREADS", str(num_threads))
 
-        def circuit():
+        def circuit(measurement):
             """A combination of the one_qubit_block and a simple PauliZ measurement applied to a
             basis state"""
             qml.BasisState(np.array(basis_state), wires=0)
             one_qubit_block(wires=0)
-            return qml.expval(qml.PauliZ(0))
+            return measurement() if callable(measurement) else measurement
 
         dev_l = lightning_dev_version(wires)
         dev_d = default_qubit_dev(wires)
 
         lightning = qml.QNode(circuit, dev_l)
         default = qml.QNode(circuit, dev_d)
 
-        lightning()
+        lightning(qml.expval(qml.PauliZ(0)))
         lightning_state = dev_l.state
 
-        default()
-        default_state = dev_d.state
+        default_state = default(qml.state)
 
         assert np.allclose(lightning_state, default_state)
         assert os.getenv("OMP_NUM_THREADS") == str(num_threads)
@@ -108,7 +107,7 @@ def test_two_qubit_circuit(
         """Test a two-qubit circuit"""
         monkeypatch.setenv("OMP_NUM_THREADS", str(num_threads))
 
-        def circuit():
+        def circuit(measurement):
             """A combination of two qubit gates with the one_qubit_block and a simple PauliZ
             measurement applied to an input basis state"""
             qml.BasisState(np.array(basis_state), wires=[0, 1])
@@ -123,19 +122,18 @@ def circuit():
             one_qubit_block(wires=1)
             qml.CRZ(0.02, wires=[0, 1])
             qml.CRot(0.2, 0.3, 0.7, wires=[0, 1])
-            return qml.expval(qml.PauliZ(0))
+            return measurement() if callable(measurement) else measurement
 
         dev_l = lightning_dev_version(wires)
         dev_d = default_qubit_dev(wires)
 
         lightning = qml.QNode(circuit, dev_l)
         default = qml.QNode(circuit, dev_d)
 
-        lightning()
+        lightning(qml.expval(qml.PauliZ(0)))
         lightning_state = dev_l.state
 
-        default()
-        default_state = dev_d.state
+        default_state = default(qml.state)
 
         assert np.allclose(lightning_state, default_state)
 
@@ -152,7 +150,7 @@ def test_three_qubit_circuit(
         """Test a three-qubit circuit"""
         monkeypatch.setenv("OMP_NUM_THREADS", str(num_threads))
 
-        def circuit():
+        def circuit(measurement):
             """A combination of two and three qubit gates with the one_qubit_block and a simple
             PauliZ measurement applied to an input basis state"""
             qml.BasisState(np.array(basis_state), wires=[0, 1, 2])
@@ -175,19 +173,18 @@ def circuit():
             qml.CRot(0.2, 0.3, 0.7, wires=[2, 1])
             qml.RZ(0.4, wires=0)
             qml.Toffoli(wires=[2, 1, 0])
-            return qml.expval(qml.PauliZ(0))
+            return measurement() if callable(measurement) else measurement
 
         dev_l = lightning_dev_version(wires)
         dev_d = default_qubit_dev(wires)
 
         lightning = qml.QNode(circuit, dev_l)
         default = qml.QNode(circuit, dev_d)
 
-        lightning()
+        lightning(qml.expval(qml.PauliZ(0)))
         lightning_state = dev_l.state
 
-        default()
-        default_state = dev_d.state
+        default_state = default(qml.state)
 
         assert np.allclose(lightning_state, default_state)
 
@@ -204,7 +201,7 @@ def test_four_qubit_circuit(
         """Test a four-qubit circuit"""
         monkeypatch.setenv("OMP_NUM_THREADS", str(num_threads))
 
-        def circuit():
+        def circuit(measurement):
             """A combination of two and three qubit gates with the one_qubit_block and a simple
             PauliZ measurement, all acting on a four qubit input basis state"""
             qml.BasisState(np.array(basis_state), wires=[0, 1, 2, 3])
@@ -232,19 +229,18 @@ def circuit():
             qml.CRot(0.2, 0.3, 0.7, wires=[2, 1])
             qml.RZ(0.4, wires=0)
             qml.Toffoli(wires=[2, 1, 0])
-            return qml.expval(qml.PauliZ(0))
+            return measurement() if callable(measurement) else measurement
 
         dev_l = lightning_dev_version(wires)
         dev_d = default_qubit_dev(wires)
 
         lightning = qml.QNode(circuit, dev_l)
         default = qml.QNode(circuit, dev_d)
 
-        lightning()
+        lightning(qml.expval(qml.PauliZ(0)))
         lightning_state = dev_l.state
 
-        default()
-        default_state = dev_d.state
+        default_state = default(qml.state)
 
         assert np.allclose(lightning_state, default_state)
 
@@ -265,23 +261,22 @@ def test_n_qubit_circuit(
         shape = qml.StronglyEntanglingLayers.shape(2, wires)
         w = np.random.uniform(high=2 * np.pi, size=shape)
 
-        def circuit():
+        def circuit(measurement):
             """Prepares the equal superposition state and then applies StronglyEntanglingLayers
             and concludes with a simple PauliZ measurement"""
             stateprep(vec, wires=range(wires))
             qml.StronglyEntanglingLayers(w, wires=range(wires))
-            return qml.expval(qml.PauliZ(0))
+            return measurement() if callable(measurement) else measurement
 
         dev_l = lightning_dev_version(wires)
         dev_d = default_qubit_dev(wires)
 
         lightning = qml.QNode(circuit, dev_l)
         default = qml.QNode(circuit, dev_d)
 
-        lightning()
+        lightning(qml.expval(qml.PauliZ(0)))
         lightning_state = dev_l.state
 
-        default()
-        default_state = dev_d.state
+        default_state = default(qml.state)
 
         assert np.allclose(lightning_state, default_state)