Adjoint of StatePrep giving unexpected phase

I am conditionally unpreparing an amplitude encoded state that I previously prepared with qml.StatePrep, by using its adjoint qml.adjoint(qml.StatePrep)). Since I start with the [1,0,0,0] state for a 2-qubit system for example, I would expect after preparing the state and then subsquently applying the adjoint of this preparation that I would be returned to the same ground state, but instead I am left with a phase offset.

  • 0 negative amplitudes to prepare, gives no offset
  • 1 negative amplitude gives an offset of pi/4
  • 2 negative amplitude gives an offset of pi/2
  • 3 negative amplitude gives an offset of 3pi/4
  • all negative amplitude gives an offset of pi

For three qubits, one negative amplitude with the rest positive gives a phase offset of pi/8.

I need to return the state to [1,0,0,0] (and in my personal application, this reversal is controlled by another entangled register, so I do not think a reset can work for me). Is there a built in way to account for this phase factor? If not I can apply an RZ gate to the first qubit in this register to account for it, but is there some built-in functionality to account for this?

import pennylane as qml

state = np.array([-0.5, 0.2, 0.3, 0.9, 0.5, 0.2, 0.3, 0.9], requires_grad=True)
state = state / np.linalg.norm(state)

test = qml.device('default.qubit', wires=3)

@qml.qnode(test)
def test_circ(state):
    qml.StatePrep(state=state,wires=[0,1,2])
    qml.adjoint(qml.StatePrep(state=state,wires=[0,1,2]))
    # optionally apply some RZ shift to correct the phase offset
    # in this specific case, it is off by pi/8
    # qml.RZ(phi=np.pi/4,wires=0)

    return qml.state()

print(test_circ(state))

The state I get instead of [1,0,0,0,0,0,0,0] is:

[ 9.23879533e-01+3.82683432e-01j -9.68859404e-18+5.12040449e-17j
  1.96261557e-16+1.27570012e-16j  1.70680150e-17+2.90657821e-17j
 -1.11022302e-16+0.00000000e+00j  9.68859404e-18-9.68859404e-18j
  3.92523115e-17-9.81307787e-18j -1.70680150e-17+1.70680150e-17j]

Also, if I had to manually adjust it with an RZ gate, would that mean I can no longer batch since the RZ angle would be dependent on the number of negative amplitudes on the prepared state?

PS: Sorry for all the edits, I made a small mistake discuss the RZ phase angle offset correction

Hi @Anthony_Smaldone,

Thanks for spotting this. Behind the scenes, StatePrep can occasionally be decomposed to MottonenStatePreparation, which is only correct up to a global phase. This is fine except for when the preparation needs to be controlled, as in your case.

Would you mind opening an issue? We’d be happy to work on a bugfix, but in the meantime you could manually place GlobalPhase (which can also be controlled on the other register).

Thank you @Tom_Bromley. I have a couple of questions. This is my GlobalPhase/RZ fix that works correctly. Since you must use GlobalPhase on all wires, an RZ gate would be better (functionally equivalent to a GlobalPhase in the case of it being applied on the ground state vector) since that can be controlled and only act on the register of interest.

@qml.qnode(test)
def test_circ(state):
    qml.StatePrep(state=state,wires=[0,1,2])
    qml.adjoint(qml.StatePrep(state=state,wires=[0,1,2]))
    #count the number of negative amplitudes in the initial state
    #to know how to adjust the global phase
    ang = np.sum(state < 0)/(state.shape[-1]) * np.pi
    
    qml.RZ(phi=ang*2)
    #qml.GlobalPhase(phi=ang)

    return qml.state()

Even still, this fix means I can no longer parallelize/batch the circuits since it would need to count the negatives in the amplitudes of each input data. This has made 1 epoch in my person model go from taking a few minutes to 7 hours. Furthermore, running each circuit and stacking them back together (I am using a TorchLayer) causes it to have differentiability issues, with the quantum weight turning to nan after one step of the optimizer.

Under what circumstances would StatePrep be decomposed to MottenStatePreparation, is this cause the backpropogation errors?

Hello Again @Tom_Bromley ,

I built a TorchLayer example that caused the quantum weight to go to nan. Is this a separate issue?

import torch
import pennylane as qml
from pennylane.qnn import TorchLayer
import numpy as np

torch.manual_seed(0)

def SimpleQuantumLayer():
    n_qubits = 2
    dev = qml.device("default.qubit", wires=n_qubits)

    def circuit(inputs, weights):
        qml.StatePrep(state=inputs,wires=[0,1])
        qml.adjoint(qml.StatePrep(state=inputs,wires=[0,1]))
        qml.RY(phi=weights[0],wires=[0])
        return qml.expval(qml.PauliZ(0))

    qlayer = qml.QNode(circuit, dev, interface="torch")
    weight_shapes = {"weights": (1, )}

    return TorchLayer(qlayer, weight_shapes)

# Create the quantum layer
quantum_layer = SimpleQuantumLayer()

# Adjusting the dimension of each vector in the batch to 4
vector_dim_adjusted = 4

# Creating a new batch of random vectors with adjusted dimensions
batch_adjusted = torch.randn(20, vector_dim_adjusted)

# L2 normalizing each vector in the adjusted batch
normalized_batch_adjusted = torch.nn.functional.normalize(batch_adjusted, p=2, dim=1)

normalized_batch_adjusted, normalized_batch_adjusted.shape

# Assuming the same scalar target for all inputs in the batch
target_value = np.cos(np.pi / 4)
target = torch.tensor(target_value, dtype=torch.float32)  # Target as a scalar



# Set up the optimizer
optimizer = torch.optim.Adam(quantum_layer.parameters(), lr=0.1)

steps = 5  # Number of optimization steps
print("Starting Parameter:", [p.item() for p in quantum_layer.parameters()])
for step in range(steps):
    total_loss = 0

    for vec in normalized_batch_adjusted:  # Loop through each vector in the batch
        optimizer.zero_grad()

        # Forward pass with the quantum layer
        output = quantum_layer(vec)
        output = torch.sqrt(output)
        prediction = output

        # Compute loss and perform backward pass
        loss = torch.nn.functional.mse_loss(prediction, target)
        total_loss += loss.item()
        loss.backward()
        optimizer.step()

    # Print average loss every few steps
    if step % 1 == 0:
        print(f"Step {step}: Average loss = {total_loss / len(normalized_batch_adjusted)}")
        print("Parameters:", [p.item() for p in quantum_layer.parameters()])

print(f"Learned angle: {list(quantum_layer.parameters())[0].item()}")

Output:

Starting Parameter: [3.118072271347046]
Step 0: Average loss = nan
Parameters: [nan]
Step 1: Average loss = nan
Parameters: [nan]
Step 2: Average loss = nan
Parameters: [nan]
Step 3: Average loss = nan
Parameters: [nan]
Step 4: Average loss = nan
Parameters: [nan]
Learned angle: nan

If you randomize the seed, it sometimes works. Is the adjoint of StatePrep not differentiable?

Hi @Anthony_Smaldone,
Thank you for your question!
We have forwarded this question to members of our technical team who will be getting back to you within a week. Feel free to post any updates to your question here in the thread in the meantime!

Hi @Anthony_Smaldone! It looks like the nan problem might be related to the earlier issue you created. Unfortunately we don’t have a fix for that right now, but this is a good datapoint that it is important!

Under what circumstances would StatePrep be decomposed to MottenStatePreparation, is this cause the backpropogation errors?

StatePrep will be decomposed whenever it is not the first operation. So one way to fix the phase issue could be to place an identity before the first StatePrep:

import numpy as np
from scipy.stats import unitary_group

state = unitary_group.rvs(2 ** 3)[:, 0]

@qml.qnode(qml.device("default.qubit"))
def test_circ(state):
    qml.Identity(0)
    qml.StatePrep(state=state,wires=[0,1,2])
    qml.adjoint(qml.StatePrep(state=state,wires=[0,1,2]))
    return qml.state()

np.round(np.real_if_close(test_circ(state)), 10)

Here I believe the phases of StatePrep and its adjoint are cancelling out, and you should be able to use that in a controlled setting too. However, this may not solve the differentiability issues you’ve been seeing. Have you tried switching to another differentiation method like diff_method="adjoint"?