PrepSelPrep template

[willjmax]

This PR introduces a template implementing qml.PrepSelPrep to perform a block encoding on a linear combination of unitaries. This is a work in progress as it does not yet support autograd or JAX. It can be used in the following way.

lcu = qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)])

dev = qml.device('default.qubit')

@qml.qnode(dev)
def circuit():
    qml.PrepSelPrep(lcu, control = 0)  
    return qml.state()

Which is equivalent to

coeffs, unitaries = lcu.terms()
normalized_coeffs = np.sqrt(coeffs) / np.linalg.norm(np.sqrt(coeffs))

@qml.qnode(dev)
def circuit():
    qml.StatePrep(normalized_coeffs, wires=0)
    qml.Select(unitaries, control=0)
    qml.adjoint(qml.StatePrep(normalized_coeffs, wires=0))

The tests in test_prepselprep.py pass, and are inspired by similar tests found in test_select.py. However, feedback on two tests would be appreciated.

1. In the assert_valid test, within _check_decomposition, there is a check to see if the operations in the decomposition match the operations in the queue. I was able to get this test to pass by adding with qml.QueuingManager.stop_recording() to the decomposition function. The decomposition function calls qml.Select and qml.StatePrep, so my best guess is that the operations were being added to the queue twice: once in qml.PrepSelPrep and again in the two subroutines. I haven't figured out how the queue works yet, so this could be totally wrong.

2. In assert_valid the _check_pytree test fails when JAX is available, but passes otherwise.

Differentiation

Here are two proposed autograd tests, each based off of the one found in test_select.py, and both fail in different ways. These are not full tests, but stripped down to the minimum required to produce the errors.

def test_autograd1():
    dev = qml.device("default.qubit")

    coeffs = pnp.array([1/2, 1/2], requires_grad=True)
    ops = [qml.Identity(1), qml.PauliZ(1)]
    lcu = qml.dot(coeffs, ops)

    @qml.qnode(dev)
    def circuit(coeffs, ops):
        lcu = qml.dot(coeffs, ops)
        qml.PrepSelPrep(lcu, control=0)
        return qml.expval(qml.Identity(1))

    qml.grad(circuit)(coeffs, ops)


def test_autograd2():
    dev = qml.device("default.qubit")

    coeffs = pnp.array([1/2, 1/2], requires_grad=True)
    ops = [qml.Identity(1), qml.PauliZ(1)]
    lcu = qml.dot(coeffs, ops)

    @qml.qnode(dev)
    def circuit(lcu):
        qml.PrepSelPrep(lcu, control=0)
        return qml.expval(qml.Identity(1))

    qml.grad(circuit)(lcu)

1. test_autograd1 fails with the error TypeError: loop of ufunc does not support argument 0 of type ArrayBox which has no callable sqrt method. This is being caused within the decomposition function by the line normalized_coeffs = np.sqrt(coeffs) / np.linalg.norm(np.sqrt(coeffs)). The issue is that coeffs is a list of autograd Arrayboxes and np.sqrt cannot be called on them. I think I should be using the functions in qml.math as they will check the types of the input and call the appropriate functions, but there is no qml.math.sqrt to call.

2. test_autograd2 passes, but with a warning since lcu is of type pennylane.ops.op_math.sum.Sum and does not have a requires_grad attribute. Attempting to force it to differentiate with qml.grad(circuit, argnum=0) results in a type error since pennylane.ops.op_math.sum.Sum is not differentiable.

Related GitHub Issues: #5739

[willjmax]

@soranjh This is the PR for the PrepSelPrep template.

[willjmax]

@soranjh The PR is ready for review.

[soranjh]

@soranjh The PR is ready for review.

Thanks @willjmax, any reason you have two modules and two classes with different names?

[willjmax]

Good catch. I renamed the module to be consistent with the name you used originally, but the old file didn't get removed. I'll get rid of it.

[soranjh]

Thanks @willjmax, left three high-level comments.

[KetpuntoG]

Hello @willjmax. Thank you for working on this template.
Here are three comments that will help you to complete the tests

[Jaybsoni]

Great work so far, I left a few comments and suggestions for you to tackle to fix up the implementation. Feel free to re-request a review once you have addressed the feedback.

Let us know if you require any clarification.

[willjmax]

@soranjh Thanks. I will have time to address these comments tomorrow.

Edit: I was able to do this tonight.

[willjmax]

@soranjh Please see my comments about your counterexample LCU.

[KetpuntoG]

Thanks @willjmax, definitely a very useful template

[Jaybsoni]

looks good to me