Repository: PennyLaneAI/pennylane
Pull request #5547 (2024-04-19 19:18:24)

Do not cast state to complex128
### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [x] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [x] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:**
The LQ new device API does not preserve the `dtype` of measurement results,

``` py
import pennylane as qml
import numpy as np
dev = qml.device("lightning.qubit", wires=2, c_dtype=np.complex64)

@qml.qnode(dev)
def circ():
    return qml.state()

>>> circ().dtype
complex128
```

The issue comes from `measurementprocess.process_state` as this method changes the specified `dtype` to the default `complex128`. 

**Description of the Change:**
Modify `StateMP` and `DensityMatrixMP` avoiding explicitly casting to `complex128`, relying on the various frameworks casting rules, by adding `0.0j`. This does not work in TensorFlow for which the current behaviour is preserved.  

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**
[sc-60855]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5547?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.67%. Comparing base [(`d03216e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/d03216e6ce5407a7f28489aa4d4b23a9c4345fac?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`3242541`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5547?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 2 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff            @@
##           master    #5547    +/-   ##
========================================
  Coverage   99.67%   99.67%            
========================================
  Files         407      407            
  Lines       38019    37754   -265     
========================================
- Hits        37895    37632   -263     
+ Misses        124      122     -2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5547?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1289ec086..aeeeedf42 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -360,6 +360,9 @@
 
 <h3>Breaking changes 💔</h3>
 
+* State measurements preserve `dtype`.
+  [(#5547)](https://github.com/PennyLaneAI/pennylane/pull/5547)
+
 * Use `SampleMP`s in the `dynamic_one_shot` transform to get back the values of the mid-circuit measurements.
   [(#5486)](https://github.com/PennyLaneAI/pennylane/pull/5486)
 
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index c31e3b94c..f0d158cf6 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -156,9 +156,10 @@ class StateMP(StateMeasurement):
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         # pylint:disable=redefined-outer-name
+        is_tf_interface = qml.math.get_deep_interface(state) == "tensorflow"
         wires = self.wires
         if not wires or wire_order == wires:
-            return qml.math.cast(state, "complex128")
+            return qml.math.cast(state, "complex128") if is_tf_interface else state + 0.0j
 
         if set(wires) != set(wire_order):
             raise WireError(
@@ -178,7 +179,7 @@ class StateMP(StateMeasurement):
         state = qml.math.reshape(state, shape)
         state = qml.math.transpose(state, desired_axes)
         state = qml.math.reshape(state, flat_shape)
-        return qml.math.cast(state, "complex128")
+        return qml.math.cast(state, "complex128") if is_tf_interface else state + 0.0j
 
 
 class DensityMatrixMP(StateMP):
@@ -211,4 +212,7 @@ class DensityMatrixMP(StateMP):
         # pylint:disable=redefined-outer-name
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
-        return qml.math.reduce_statevector(state, indices=mapped_wires)
+        kwargs = {"indices": mapped_wires, "c_dtype": "complex128"}
+        if not qml.math.is_abstract(state) and qml.math.any(qml.math.iscomplex(state)):
+            kwargs["c_dtype"] = state.dtype
+        return qml.math.reduce_statevector(state, **kwargs)


Test patch:
diff --git a/tests/devices/test_lightning_qubit.py b/tests/devices/test_lightning_qubit.py
index 9493382d4..fe97c9c8b 100644
--- a/tests/devices/test_lightning_qubit.py
+++ b/tests/devices/test_lightning_qubit.py
@@ -98,7 +98,7 @@ class TestDtypePreserved:
     @pytest.mark.parametrize(
         "c_dtype",
         [
-            pytest.param(np.complex64, marks=pytest.mark.xfail(reason="dtype not preserved")),
+            np.complex64,
             np.complex128,
         ],
     )
@@ -108,18 +108,10 @@ class TestDtypePreserved:
             qml.state(),
             qml.density_matrix(wires=[1]),
             qml.density_matrix(wires=[2, 0]),
-            pytest.param(
-                qml.expval(qml.PauliY(0)), marks=pytest.mark.xfail(reason="incorrect type")
-            ),
-            pytest.param(qml.var(qml.PauliY(0)), marks=pytest.mark.xfail(reason="incorrect type")),
-            pytest.param(
-                qml.probs(wires=[1]),
-                marks=pytest.mark.skip(reason="measurement passes with complex64 but xfail strict"),
-            ),
-            pytest.param(
-                qml.probs(wires=[0, 2]),
-                marks=pytest.mark.skip(reason="measurement passes with complex64 but xfail strict"),
-            ),
+            qml.expval(qml.PauliY(0)),
+            qml.var(qml.PauliY(0)),
+            qml.probs(wires=[1]),
+            qml.probs(wires=[0, 2]),
         ],
     )
     def test_dtype(self, c_dtype, measurement):
@@ -139,4 +131,7 @@ class TestDtypePreserved:
             expected_dtype = c_dtype
         else:
             expected_dtype = np.float64 if c_dtype == np.complex128 else np.float32
-        assert res.dtype == expected_dtype
+        if isinstance(res, np.ndarray):
+            assert res.dtype == expected_dtype
+        else:
+            assert isinstance(res, float)
diff --git a/tests/gradients/core/test_pulse_gradient.py b/tests/gradients/core/test_pulse_gradient.py
index b95a36d67..04d6b7e35 100644
--- a/tests/gradients/core/test_pulse_gradient.py
+++ b/tests/gradients/core/test_pulse_gradient.py
@@ -770,7 +770,7 @@ class TestStochPulseGradErrors:
         x = [0.4, 0.2]
         params = [jnp.array(0.14)]
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
-        op = qml.evolve(ham_single_q_const)(params, 0.1)
+        op = qml.evolve(ham_single_q_const)(params, 0.7)
         tape = qml.tape.QuantumScript(
             [qml.RX(x, 0), op], [qml.expval(qml.PauliZ(0))], trainable_params=[1]
         )