Repository: PennyLaneAI/pennylane
Pull request #4811 (2023-11-08 22:41:04)

ensure the correct interface is returned in expand_state
**Context:**
In the past, QNodes would return pennylane tensors by default if there was no specific interface set. However, I introduced some funniness when I added `qml.math.pad` [here](https://github.com/PennyLaneAI/pennylane/blob/4eecda939f862d62662751e791ba58bb0fc13508/pennylane/devices/qubit/simulate.py#L90). For some reason, when given a pennylane tensor, it would return a vanilla numpy array.

**Description of the Change:**
Provide the `like` parameter to `qml.math.pad` when expanding the statevector over all wires.

**Benefits:**
Consistency with past versions of PennyLane

**Possible Drawbacks:**
What in tarnation. I can't make any sense of interfaces anymore.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4811?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:
> Comparison is base [(`4eecda9`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/4eecda939f862d62662751e791ba58bb0fc13508?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`96c3cd9`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4811?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.

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


```diff
@@            Coverage Diff             @@
##           master    #4811      +/-   ##
==========================================
- Coverage   99.64%   99.63%   -0.01%     
==========================================
  Files         381      381              
  Lines       34334    34075     -259     
==========================================
- Hits        34212    33952     -260     
- Misses        122      123       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4811?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/qubit/simulate.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4811?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvc2ltdWxhdGUucHk=) | `100.00% <100.00%> (ø)` | |

... and [43 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4811/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4811?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 7c5fd3ece..0fa5a007d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -96,6 +96,10 @@
 * `transpile` can now handle measurements that are broadcasted onto all wires.
   [(#4793)](https://github.com/PennyLaneAI/pennylane/pull/4793)
 
+* Parametrized circuits whose operators do not act on all wires return pennylane tensors as
+  expected, instead of numpy arrays.
+  [(#4811)](https://github.com/PennyLaneAI/pennylane/pull/4811)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 22b67e474..7a4bb6c57 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -75,8 +75,9 @@ def expand_state_over_wires(state, state_wires, all_wires, is_state_batched):
     Returns:
         TensorLike: The state in the new desired size and order
     """
+    interface = qml.math.get_interface(state)
     pad_width = 2 ** len(all_wires) - 2 ** len(state_wires)
-    pad = (pad_width, 0) if qml.math.get_interface(state) == "torch" else (0, pad_width)
+    pad = (pad_width, 0) if interface == "torch" else (0, pad_width)
     shape = (2,) * len(all_wires)
     if is_state_batched:
         pad = ((0, 0), pad)
@@ -87,7 +88,7 @@ def expand_state_over_wires(state, state_wires, all_wires, is_state_batched):
         pad = (pad,)
         state = qml.math.flatten(state)
 
-    state = qml.math.pad(state, pad, mode="constant")
+    state = qml.math.pad(state, pad, mode="constant", like=interface)
     state = qml.math.reshape(state, shape)
 
     # re-order


Test patch:
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index eb1877a21..3b0471a98 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -189,6 +189,16 @@ class TestBasicCircuit:
         assert qml.math.get_interface(res) == "jax"
         assert qml.math.allclose(res, -1)
 
+    def test_expand_state_keeps_autograd_interface(self):
+        """Test that expand_state doesn't convert autograd to numpy."""
+
+        @qml.qnode(qml.device("default.qubit", wires=2))
+        def circuit(x):
+            qml.RX(x, 0)
+            return qml.probs(wires=[0, 1])
+
+        assert qml.math.get_interface(circuit(1.5)) == "autograd"
+
 
 class TestBroadcasting:
     """Test that simulate works with broadcasted parameters"""