search

PennyLaneAI / pennylane

PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
https://pennylane.ai
Apache License 2.0
2.34k stars 602 forks source link

Printing or storing intermediate values inside a cost function produces ArrayBox objects #1001

Closed antalszava closed 3 years ago

antalszava commented 3 years ago

Issue description

When using with the step_and_cost function to optimize a cost function that involves computing probabilities with qml.probs, the output values of the QNode are stored as ArrayBox objects.

import pennylane as qml
from pennylane import numpy as np

wires = 6
dev = qml.device("default.qubit", wires=wires, shots=1000, analytic=False)

def global_cost_simple(rotations):
    for i in range(wires):
        qml.RX(rotations[0][i], wires=i)
    return qml.probs(wires=range(wires))

global_circuit = qml.QNode(global_cost_simple, dev)

def cost_global(rotations):
    probs = global_circuit(rotations)
    print(probs)   # <----- Outputs Autograd ArrayBox objects
    return 1 - probs[0]

params_global = np.array([[3.] * len(range(wires))])
opt = qml.GradientDescentOptimizer(stepsize=0.2)
steps = 5

for i in range(steps):
    params_global, c = opt.step_and_cost(cost_global, params_global)

Description of the issue - include code snippets and screenshots here if relevant. You may use the following template below

Additional information

A solution can be to store the output of dev.probability(wires=dev.wires) (dev is the PennyLane device) which does involves computations using the statevector/samples without executing the QNode again.

josh146 commented 3 years ago

@antalszava, this might be expected behaviour (unfortunately), and is due to how Autograd works.

The new step_and_cost method takes advantage of the fact that when you backpropagate with Autograd via qml.grad(), several things happen:

So step_and_cost allows you to step the parameters and retrieve the cost function with minimal overhead, but only the output of the function will not be an ArrayBox. Everything intermediate (as is the case of the print statement) will be an ArrayBox.

josh146 commented 3 years ago

For example, consider the following pure Autograd code:

from autograd import numpy as np

from autograd.core import make_vjp
from autograd.wrap_util import unary_to_nary
from autograd.extend import vspace

@unary_to_nary
def grad(fun, x):
    """This function is a replica of autograd.grad
    modified to also return the function output"""
    vjp, fn_output = _make_vjp(fun, x)

    if not vspace(fn_output).size == 1:
        raise TypeError("Grad only applies to real scalar-output functions. Try jacobian or elementwise_grad.")

    grad_value = vjp(vspace(fn_output).ones())

    # the following return statement is modified;
    # the standard autograd.grad only returns grad_value here
    return grad_value, fn_output

def function(x):
    intermediate = x ** 2
    print("This intermediate value will be an ArrayBox:", intermediate)
    return 1 - intermediate

grad_fn = grad(function, argnum=0)
grad_value, cost_value = grad_fn(0.5)

print("Gradient value is not an ArrayBox:", grad_value)
print("Function output is not an ArrayBox:", cost_value)

This has output:

This intermediate value will be an ArrayBox: Autograd ArrayBox with value 0.25
Gradient value is not an ArrayBox: -1.0
Function output is not an ArrayBox: 0.75
Francesco-Benfenati commented 3 years ago

So using step_and_cost I won't be able to have the probabilities as numbers, and the two possible solutions would be accessing dev.probability(wires=dev.wires) or splitting step_and_cost into step and cost, both of which would require twice the number of executions of the qnode, right? The most important thing for me would be to avoid measuring the same circuit twice, that's why I thought that I should use step_and_cost.

Would this be the case even with real QPUs? Because for example I think that IBMQ is outputting the probabilities of all the states as numbers, so it should be possible to print them and use them to calculate the cost without repeating the experiment.

josh146 commented 3 years ago

Hi @Francesco-Benfenati,

splitting step_and_cost into step and cost, both of which would require twice the number of executions of the qnode, right?

Yes, exactly right - step_and_cost computes the gradient and extracts the corresponding forward pass value at the same time. Calling step and cost separately would instead result in two executions.

accessing dev.probability(wires=dev.wires)

This approach actually doesn't result in an extra evaluation. The device always stores the result of the last computation, so calling dev.probability just extracts the last execution probabilities 🙂

Would this be the case even with real QPUs? Because for example I think that IBMQ is outputting the probabilities of all the states as numbers, so it should be possible to print them and use them to calculate the cost without repeating the experiment.

Unfortunately yes, the issue currently stems from how Autograd works, which is the default backend PennyLane uses for autodifferentiation. If you compute the gradient in Autograd, the cost function value is always computed implicitly at the same time, and all intermediate values of the computation will be Autograd ArrayBox objects.

I think your solution of using _value is actually quite a neat one:

from autograd import numpy as np

from autograd.core import make_vjp
from autograd.wrap_util import unary_to_nary
from autograd.extend import vspace

@unary_to_nary
def grad(fun, x):
    """This function is a replica of autograd.grad
    modified to also return the function output"""
    vjp, fn_output = make_vjp(fun, x)

    if not vspace(fn_output).size == 1:
        raise TypeError(
            "Grad only applies to real scalar-output functions. "
            "Try jacobian, elementwise_grad or holomorphic_grad."
        )

    grad_value = vjp(vspace(fn_output).ones())
    return grad_value, fn_output

def function(x):
    intermediate = x ** 2
    print("This intermediate value will be an ArrayBox:", intermediate)
    print("Calling ._value extracts the NumPy array:", intermediate._value)
    return 1 - intermediate

grad_fn = grad(function, argnum=0)
grad_value, cost_value = grad_fn(0.5)

print("Gradient value is not an ArrayBox:", grad_value)
print("Function output is not an ArrayBox:", cost_value)

This gives output

This intermediate value will be an ArrayBox: Autograd ArrayBox with value 0.25
Calling ._value extracts the NumPy array: 0.25
Gradient value is not an ArrayBox: -1.0
Function output is not an ArrayBox: 0.75

Alternatively, if you switch to using the TensorFlow or PyTorch interfaces, this should no longer be an issue (since they do not use ArrayBox objects).

Francesco-Benfenati commented 3 years ago

Thank you @josh146 and @antalszava ! I'll experiment with that