Repository: PennyLaneAI/pennylane

Issue:
StatePrep sometimes fails for TorchLayers
### Expected behavior

When using StatePrep for statevectors with torch gradients, the `StatePrep` function sometimes causes PyTorch to throw an error. It should be able to prepare the state on the specified subsystem of qubits, and allow for other gates to be executed while maintaining the input feature (state vector) gradients.

### Actual behavior

However, this throws an error with PyTorch:

```
RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.
```

### Additional information

If you change the state preparation to `wires=[0,1,2]`. There is no error. If the `RY` gate is removed, there is no error. If you detach the input features from the torch gradients, there is no error.

### Source code

```shell
import torch
import torch.nn as nn
import pennylane as qml
import pennylane.numpy as np

def QuantumLayer():
    n_qubits = 4
    dev = qml.device("default.qubit", wires=n_qubits)

    def _circuit(inputs, weights):
        qml.StatePrep(inputs, wires=[1,2,3])
        qml.RY(phi=weights, wires=[0])
        return qml.expval(qml.PauliZ(wires=0))
    
    qlayer = qml.QNode(_circuit, dev, interface="torch")
    weight_shapes = {"weights": (1)}
    return qml.qnn.TorchLayer(qlayer, weight_shapes)

class SimpleQuantumModel(nn.Module):
    def __init__(self):
        super(SimpleQuantumModel, self).__init__()
        self.quantum_layer = QuantumLayer()
    def forward(self, x):
        return self.quantum_layer(x)

model = SimpleQuantumModel()
features = torch.tensor(np.array([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]), requires_grad=True)

print(model(features))
```


### Tracebacks

```shell
Traceback (most recent call last):
  File "troubleshooting_forum.py", line 31, in <module>
    print(model(features))
  File "<python_env_path>/lib/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "<python_env_path>/lib/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "troubleshooting_forum.py", line 26, in forward
    return self.quantum_layer(x)
  File "<python_env_path>/lib/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "<python_env_path>/lib/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "<python_env_path>/lib/site-packages/pennylane/qnn/torch.py", line 402, in forward
    results = self._evaluate_qnode(inputs)
  File "<python_env_path>/lib/site-packages/pennylane/qnn/torch.py", line 423, in _evaluate_qnode
    res = self.qnode(**kwargs)
  File "<python_env_path>/lib/site-packages/pennylane/qnode.py", line 1027, in __call__
    res = qml.execute(
  File "<python_env_path>/lib/site-packages/pennylane/interfaces/execution.py", line 616, in execute
    results = inner_execute(tapes)
  File "<python_env_path>/lib/site-packages/pennylane/interfaces/execution.py", line 249, in inner_execute
    return cached_device_execution(tapes)
  File "<python_env_path>/lib/site-packages/pennylane/interfaces/execution.py", line 371, in wrapper
    res = list(fn(tuple(execution_tapes.values()), **kwargs))
  File "<python_env_path>/lib/site-packages/pennylane/devices/default_qubit.py", line 474, in execute
    results = tuple(
  File "<python_env_path>/lib/site-packages/pennylane/devices/default_qubit.py", line 475, in <genexpr>
    simulate(
  File "<python_env_path>/lib/site-packages/pennylane/devices/qubit/simulate.py", line 269, in simulate
    state, is_state_batched = get_final_state(circuit, debugger=debugger, interface=interface)
  File "<python_env_path>/lib/site-packages/pennylane/devices/qubit/simulate.py", line 156, in get_final_state
    state = create_initial_state(sorted(circuit.op_wires), prep, like=INTERFACE_TO_LIKE[interface])
  File "<python_env_path>/lib/site-packages/pennylane/devices/qubit/initialize_state.py", line 45, in create_initial_state
    return qml.math.asarray(prep_operation.state_vector(wire_order=list(wires)), like=like)
  File "<python_env_path>/lib/site-packages/pennylane/ops/qubit/state_preparation.py", line 230, in state_vector
    ket[indices] = op_vector
  File "<python_env_path>/lib/site-packages/torch/_tensor.py", line 1032, in __array__
    return self.numpy().astype(dtype, copy=False)
RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.
```


### System information

```shell
Name: PennyLane
Version: 0.33.1
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /usr/local/lib/python3.10/dist-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           Linux-5.15.120+-x86_64-with-glibc2.35
Python version:          3.10.12
Numpy version:           1.23.5
Scipy version:           1.11.3
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:
Thanks for posting this issue @anthonysmaldone . 

This does indeed seem to be a bug. It looks like we aren't expanding the state in a backprop-compatible way. 

We'll try and get a fix in as soon as we can, but until then I recommend setting `diff_method="adjoint"` in the QNode.  
> Thanks for posting this issue @anthonysmaldone .
> 
> This does indeed seem to be a bug. It looks like we aren't expanding the state in a backprop-compatible way.
> 
> We'll try and get a fix in as soon as we can, but until then I recommend setting `diff_method="adjoint"` in the QNode.

Thank you for looking into this. While the `diff_method="adjoint"` worked on the toy example I provided for a single circuit, it still failed to calculate gradients with my actual model, giving me the error:

```
RuntimeError: Function ExecuteTapesBackward returned an invalid gradient at index 0 - got [] but expected shape compatible with [8]
```

Both `adjoint` and the `parameter-shift` method do not seem to handle batch inputs properly either, giving an error in the provided model as:

```
RuntimeError: shape '[2]' is invalid for input of size 1
```

Thus my question becomes:

Are there any other ways to amplitude encode data where the input features are differentiable? I tried `MottonenStatePreparation` but had some issues calculating gradients as well. 
I am not sure if this is a symptom of the same issue, but I noticed that if you swap the order of the state preparation and the `RY` gate:
```
def _circuit(inputs, weights):
        qml.RY(phi=weights, wires=[0])
        qml.StatePrep(inputs, wires=[1,2,3])
        return qml.expval(qml.PauliZ(wires=0))
```

 this resolves the numpy/torch gradient error, but breaks the batching mechanism:
 
 ```
RuntimeError: shape '[2]' is invalid for input of size 1
```

While in theory one could run them sequentially and then stack the outputs, I noticed this results in `nan` values in the gradients after being fed to an optimizer.
So if a `StatePrep` is the first operation in a circuit, it is used to initialize the state vector. When it is not the first operation in the circuit, it is decomposed via `MottonenStatePreparation`.  Unfortunately, `MottonenStatePreparation` doesn't support batching.
@albi3ro Thanks for the clarification. After playing with using `adjoint` instead, I was still unable to implement it in my model. I have created another toy example to replicate my problem with the `adjoint` differentiation method. If I have a model where I run two circuits, where I feed the measurements of the first into the second as features (which is required for my personal program), it seems I am unable to update my quantum gradients.

Here is a simple example to illustrate this.

This trivial model:

1. Takes in a 3 qubit state and encodes it with `StatePrep`
2. Applies an arbitrary parametric `RY` gate
3. Measures all the qubits, and uses the normalized `expval`s as input to a second quantum circuit
4. The model updates with `SGD`

```
import torch
import torch.nn as nn
import pennylane as qml
import pennylane.numpy as np

def QuantumLayer(qubits):
    dev = qml.device("default.qubit", wires=qubits)

    def _circuit(inputs, weights):
        qml.StatePrep(inputs, wires=range(1, qubits))
        qml.RY(weights, wires=0)

        if qubits == 4:
            return [qml.expval(qml.PauliZ(wires=i)) for i in range(qubits)]
        else:
            return qml.expval(qml.PauliZ(wires=2))

    qlayer = qml.QNode(_circuit, dev, interface="torch", diff_method="adjoint")
    weight_shapes = {"weights": (1,)}
    return qml.qnn.TorchLayer(qlayer, weight_shapes)

class SimpleQuantumModel(nn.Module):
    def __init__(self):
        super(SimpleQuantumModel, self).__init__()
        self.quantum_layer = QuantumLayer(4)
        self.quantum_layer2 = QuantumLayer(3)
    
    def forward(self, x):
        x = self.quantum_layer(x)
        x = torch.nn.functional.normalize(x, p=2, dim=0)
        x = self.quantum_layer2(x)
        return x

# MSE Cost
def cost_fn(pred, target):
    return torch.mean((pred - target)**2)

# Initialize the model
model = SimpleQuantumModel()

# Define the optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Training data and expectation value label to learn
data = [
    (torch.tensor(np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), requires_grad=True), torch.tensor([1.0])),
    (torch.tensor(np.array([0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), requires_grad=True), torch.tensor([-1.0]))
]

# Training loop
epochs = 100
for epoch in range(epochs):
    total_loss = 0
    for x, y in data:
        optimizer.zero_grad()
        prediction = model(x)
        loss = cost_fn(prediction, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch}, Loss {total_loss/len(data)}")
```

This completes the forward pass, but fails to update the weights.

```
RuntimeError: Function ExecuteTapesBackward returned an invalid gradient at index 0 - got [] but expected shape compatible with [4]
```
You are right. Sorry about that suggestion; I should have realized that wouldn't work.

Adjoint differentiation actually can't handle differentiating state prep operations.  

Here's a more minimal example of the problem.

```
import torch

@qml.qnode(qml.device('default.qubit'), diff_method="adjoint")
def circuit(state):
    qml.StatePrep(state, wires=(0,))
    return qml.expval(qml.PauliZ(0))

state = torch.tensor([1.0, 0.0], requires_grad=True)

y = circuit(state)
y.backward()
```

Two bugs 🐛 in one :(

Resolving patch:
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -217,30 +217,17 @@ def state_vector(self, wire_order=None):
         if not wire_order.contains_wires(self.wires):
             raise WireError(f"Custom wire_order must contain all {self.name} wires")
 
-        num_total_wires = len(wire_order)
-        indices = tuple(
-            [Ellipsis] + [slice(None)] * num_op_wires + [0] * (num_total_wires - num_op_wires)
-        )
-        ket_shape = [2] * num_total_wires
+        # add zeros for each wire that isn't being set
+        extra_wires = Wires(set(wire_order) - set(self.wires))
+        for _ in extra_wires:
+            op_vector = math.stack([op_vector, math.zeros_like(op_vector)], axis=-1)
+
+        # transpose from operator wire order to provided wire order
+        current_wires = self.wires + extra_wires
+        transpose_axes = [current_wires.index(w) for w in wire_order]
         if self.batch_size:
-            # Add broadcasted dimension to the shape of the state vector
-            ket_shape = [self.batch_size] + ket_shape
-
-        ket = np.zeros(ket_shape, dtype=np.complex128)
-        ket[indices] = op_vector
-
-        # unless wire_order is [*self.wires, *rest_of_wire_order], need to rearrange
-        if self.wires != wire_order[:num_op_wires]:
-            current_order = self.wires + list(Wires.unique_wires([wire_order, self.wires]))
-            desired_order = [current_order.index(w) for w in wire_order]
-            if self.batch_size:
-                # If the operation is broadcasted, the desired order must include the batch dimension
-                # as the first dimension.
-                desired_order = [0] + [d + 1 for d in desired_order]
-
-            ket = ket.transpose(desired_order)
-
-        return math.convert_like(ket, op_vector)
+            transpose_axes = [0] + [a + 1 for a in transpose_axes]
+        return math.transpose(op_vector, transpose_axes)
 
 
 # pylint: disable=missing-class-docstring


Test patch:
diff --git a/tests/ops/qubit/test_state_prep.py b/tests/ops/qubit/test_state_prep.py
--- a/tests/ops/qubit/test_state_prep.py
+++ b/tests/ops/qubit/test_state_prep.py
@@ -157,7 +157,7 @@ def test_StatePrep_reordering_broadcasted(self):
         assert np.array_equal(ket, expected)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "torch", "tensorflow"])
     def test_StatePrep_state_vector_preserves_parameter_type(self, interface):
         """Tests that given an array of some type, the resulting state vector is also that type."""
         qsv_op = qml.StatePrep(qml.math.array([0, 0, 0, 1], like=interface), wires=[1, 2])
@@ -165,7 +165,7 @@ def test_StatePrep_state_vector_preserves_parameter_type(self, interface):
         assert qml.math.get_interface(qsv_op.state_vector(wire_order=[0, 1, 2])) == interface
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "torch", "tensorflow"])
     def test_StatePrep_state_vector_preserves_parameter_type_broadcasted(self, interface):
         """Tests that given an array of some type, the resulting state vector is also that type."""
         qsv_op = qml.StatePrep(
@@ -191,6 +191,104 @@ def test_StatePrep_wrong_param_size_fails(self):
         with pytest.raises(ValueError, match="State vector must have shape"):
             _ = qml.StatePrep([0, 1], wires=[0, 1])
 
+    @pytest.mark.torch
+    def test_StatePrep_torch_differentiable(self):
+        """Test that StatePrep works with torch."""
+        import torch
+
+        def QuantumLayer():
+            @qml.qnode(qml.device("default.qubit"), interface="torch")
+            def qlayer(inputs, weights):
+                qml.StatePrep(inputs, wires=[1, 2, 3])
+                qml.RY(phi=weights, wires=[0])
+                return qml.expval(qml.PauliZ(wires=0))
+
+            weight_shapes = {"weights": (1)}
+            return qml.qnn.TorchLayer(qlayer, weight_shapes)
+
+        class SimpleQuantumModel(torch.nn.Module):  # pylint:disable=too-few-public-methods
+            def __init__(self):
+                super().__init__()
+                self.quantum_layer = QuantumLayer()
+
+            def forward(self, x):
+                return self.quantum_layer(x)
+
+        model = SimpleQuantumModel()
+        features = torch.tensor(
+            [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]],
+            requires_grad=True,
+        )
+        result = model(features)
+        assert qml.math.get_interface(result) == "torch"
+        assert qml.math.shape(result) == (2,)
+
+    def test_StatePrep_backprop_autograd(self):
+        """Test backprop with autograd"""
+
+        @qml.qnode(qml.device("default.qubit"), diff_method="backprop")
+        def circuit(state):
+            qml.StatePrep(state, wires=(0,))
+            qml.S(1)
+            return qml.expval(qml.PauliZ(0))
+
+        state = qml.numpy.array([1.0, 0.0])
+        grad = qml.jacobian(circuit)(state)
+        assert np.array_equal(grad, [2.0, 0.0])
+
+    @pytest.mark.torch
+    def test_StatePrep_backprop_torch(self):
+        """Test backprop with torch, getting state.grad"""
+        import torch
+
+        @qml.qnode(qml.device("default.qubit"), diff_method="backprop")
+        def circuit(state):
+            qml.StatePrep(state, wires=(0,))
+            qml.S(1)
+            return qml.expval(qml.PauliZ(0))
+
+        state = torch.tensor([1.0, 0.0], requires_grad=True)
+        res = circuit(state)
+        res.backward()
+        grad = state.grad
+        assert qml.math.get_interface(grad) == "torch"
+        assert np.array_equal(grad, [2.0, 0.0])
+
+    @pytest.mark.jax
+    def test_StatePrep_backprop_jax(self):
+        """Test backprop with jax"""
+        import jax
+
+        @qml.qnode(qml.device("default.qubit"), diff_method="backprop")
+        def circuit(state):
+            qml.StatePrep(state, wires=(0,))
+            qml.S(1)
+            return qml.expval(qml.PauliZ(0))
+
+        state = jax.numpy.array([1.0, 0.0])
+        grad = jax.jacobian(circuit)(state)
+        assert qml.math.get_interface(grad) == "jax"
+        assert np.array_equal(grad, [2.0, 0.0])
+
+    @pytest.mark.tf
+    def test_StatePrep_backprop_tf(self):
+        """Test backprop with tf"""
+        import tensorflow as tf
+
+        @qml.qnode(qml.device("default.qubit"), diff_method="backprop")
+        def circuit(state):
+            qml.StatePrep(state, wires=(0,))
+            qml.S(1)
+            return qml.expval(qml.PauliZ(0))
+
+        state = tf.Variable([1.0, 0.0])
+        with tf.GradientTape() as tape:
+            res = circuit(state)
+
+        grad = tape.jacobian(res, state)
+        assert qml.math.get_interface(grad) == "tensorflow"
+        assert np.array_equal(grad, [2.0, 0.0])
+
     @pytest.mark.parametrize(
         "num_wires,wire_order,one_position",
         [
@@ -239,7 +337,7 @@ def test_BasisState_state_vector_computed(self, state, device_wires, op_wires):
         assert not np.any(basis_state)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "torch", "tensorflow"])
     @pytest.mark.parametrize("dtype_like", [0, 0.0])
     def test_BasisState_state_vector_preserves_parameter_type(self, interface, dtype_like):
         """Tests that given an array of some type, the resulting state_vector is also that type."""