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.33k stars 598 forks source link

[BUG] various problems with QubitUnitary #5308

Closed cvjjm closed 6 months ago

cvjjm commented 7 months ago

Expected behavior

QubitUnitary should be correctly executed for all unitary input matrices, should be (at least) finite difference differentiable, and parameter-shift differentiable via decomposition.

Actual behavior

The minimal non working example code reveals several problems:

I provide the output of the minimal non-working code and its output below.

Additional information

No response

Source code

import pennylane as qml
from pennylane import numpy as np

def make_unitary(theta1):
    generator = theta1 / 2 * (
        np.cos(0.2)
        / 2
        * (np.array([[0, 0, 0, 0], [0, 0, 2, 0], [0, 2, 0, 0], [0, 0, 0, 0]]))
        + np.sin(0.2)
        / 2
        * (np.array([[0, 0, 0, 0], [0, 0, 2j, 0], [0, -2j, 0, 0], [0, 0, 0, 0]]))
    )

    def expm(X):
        d, U = np.linalg.eigh(-1.0j * X)
        V = np.dot(U, np.dot(np.diag(np.exp(1.0j * d)), np.conj(U).T))
        return V

    mat = expm(-1j * generator)

    assert np.allclose(np.dot(np.transpose(np.conj(mat)), mat), np.eye(len(mat))), "mat is not unitary"

    return mat

in_mat = make_unitary(np.pi/2)
print("in_mat")
print(in_mat.round(3))

dev0 = qml.device("default.qubit", 2)
@qml.qnode(dev0, diff_method="finite-diff")
def foo():
    qml.QubitUnitary(in_mat, wires=[0, 1])
    return qml.state()

out_mat = qml.matrix(foo, wire_order=[0, 1])()
print("out_mat")
print(out_mat.round(3))
global_phase = np.angle(out_mat[0, 0] / in_mat[0, 0])
print("global_phase", global_phase)
out_mat *= np.exp(-1j * global_phase)
print("out_mat with phase corrected")
print(out_mat.round(3))
print("difference should be zero but it is not:")
print(np.round(in_mat - out_mat, 3))
print("circuit:")
print(qml.draw(foo, wire_order=[0, 1], expansion_strategy="device")())
print("tape looks correct:")
print(foo.tape.operations)
print("state returned by qnode should be |0> but it is not:")
print(foo())

@qml.qnode(dev0, diff_method="finite-diff", expansion_strategy="device")
def bar(theta):
    qml.QubitUnitary(make_unitary(theta), wires=[0, 1])
    return qml.expval(qml.PauliZ(0))

try:
    print("jac1", qml.jacobian(bar)(np.array(np.pi/2, requires_grad=True)))
except Exception as exception:
    jac2 = qml.jacobian(bar)(1e-3 * np.array(np.random.randn(1), requires_grad=True))
    print("jac2", jac2)
    if np.isnan(jac2):
        raise Exception("jac2 should not be nan") from exception

Tracebacks

in_mat
[[ 1.   +0.j     0.   +0.j     0.   +0.j     0.   +0.j   ]
 [ 0.   +0.j     0.707-0.j     0.14 -0.693j  0.   +0.j   ]
 [ 0.   +0.j    -0.14 -0.693j  0.707+0.j     0.   +0.j   ]
 [ 0.   +0.j     0.   +0.j     0.   +0.j     1.   +0.j   ]]
out_mat
[[ 0.056+0.135j  0.752-0.404j  0.339-0.102j  0.135+0.327j]
 [-0.246-0.817j -0.056-0.135j  0.347-0.068j -0.339+0.102j]
 [-0.312+0.167j  0.293-0.197j -0.056-0.135j -0.404-0.752j]
 [ 0.327-0.135j  0.167+0.312j -0.246-0.817j -0.135+0.056j]]
global_phase 1.1780971988894835
out_mat with phase corrected
[[ 0.146+0.j    -0.085-0.849j  0.035-0.352j  0.354+0.j   ]
 [-0.849-0.085j -0.146-0.j     0.07 -0.347j -0.035+0.352j]
 [ 0.035+0.352j -0.07 -0.347j -0.146-0.j    -0.849+0.085j]
 [-0.   -0.354j  0.352-0.035j -0.849-0.085j -0.   +0.146j]]
difference should be zero but it is not:
[[ 0.854-0.j     0.085+0.849j -0.035+0.352j -0.354-0.j   ]
 [ 0.849+0.085j  0.854+0.j     0.07 -0.347j  0.035-0.352j]
 [-0.035-0.352j -0.07 -0.347j  0.854+0.j     0.849-0.085j]
 [ 0.   +0.354j -0.352+0.035j  0.849+0.085j  1.   -0.146j]]
circuit:
0: ─╭U(M0)─┤ ╭State
1: ─╰U(M0)─┤ ╰State

M0 = 
[[ 1.        +0.00000000e+00j  0.        +0.00000000e+00j
   0.        +0.00000000e+00j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j  0.70710678-2.29934717e-17j
   0.14048043-6.93011723e-01j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j -0.14048043-6.93011723e-01j
   0.70710678+0.00000000e+00j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j  0.        +0.00000000e+00j
   0.        +0.00000000e+00j  1.        +0.00000000e+00j]]
tape looks correct:
[QubitUnitary(tensor([[ 1.        +0.00000000e+00j,  0.        +0.00000000e+00j,
          0.        +0.00000000e+00j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j,  0.70710678-2.29934717e-17j,
          0.14048043-6.93011723e-01j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j, -0.14048043-6.93011723e-01j,
          0.70710678+0.00000000e+00j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j,  0.        +0.00000000e+00j,
          0.        +0.00000000e+00j,  1.        +0.00000000e+00j]], requires_grad=True), wires=[0, 1])]
state returned by qnode should be |0> but it is not:
[ 0.0560427 +0.13529903j -0.24628222-0.81725055j -0.31150153+0.16723276j
  0.32664072-0.13529903j]
/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/numpy/numpy_wrapper.py:156: ComplexWarning: Casting complex values to real discards the imaginary part
  return A.astype(dtype, order, casting, subok, copy)
/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/numpy/tensor.py:155: RuntimeWarning: invalid value encountered in power
  res = super().__array_ufunc__(ufunc, method, *args, **kwargs)
/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/numpy/linalg.py:180: RuntimeWarning: divide by zero encountered in divide
  F = off_diag / (T(w_repeated) - w_repeated + anp.eye(N))
jac2 [nan]
Traceback (most recent call last):
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/core.py", line 31, in __init__
    vjpmaker = primitive_vjps[fun]
KeyError: <function primitive.<locals>.f_wrapped at 0x7fbc8d724040>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/cvjjm/src/covqcstack/hfak/covvqetools/qubit_unitary_test.py", line 59, in <module>
    print("jac1", qml.jacobian(bar)(np.array(np.pi/2, requires_grad=True)))
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/_grad.py", line 455, in _jacobian_function
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/_grad.py", line 455, in <genexpr>
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/differential_operators.py", line 60, in jacobian
    vjp, ans = _make_vjp(fun, x)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/core.py", line 10, in make_vjp
    end_value, end_node =  trace(start_node, fun, x)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/tracer.py", line 10, in trace
    end_box = fun(start_box)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/wrap_util.py", line 15, in unary_f
    return fun(*subargs, **kwargs)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/qnode.py", line 1039, in __call__
    res = qml.execute(
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/interfaces/execution.py", line 632, in execute
    tapes, post_processing = transform_program(tapes)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/transforms/core/transform_program.py", line 435, in __call__
    new_tapes, fn = transform(tape, *targs, **tkwargs)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/gradients/finite_difference.py", line 191, in _expand_transform_finite_diff
    expanded_tape = expand_invalid_trainable(tape)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/transforms/tape_expand.py", line 102, in expand_fn
    tape = tape.expand(depth=depth, stop_at=stop_at)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/tape/qscript.py", line 946, in expand
    new_script = qml.tape.tape.expand_tape(
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/tape/tape.py", line 203, in expand_tape
    obj = QuantumScript(obj.decomposition(), _update=False)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/operation.py", line 1261, in decomposition
    return self.compute_decomposition(
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/ops/qubit/matrix_ops.py", line 222, in compute_decomposition
    return qml.ops.two_qubit_decomposition(U, Wires(wires))
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/ops/op_math/decompositions/two_qubit_unitary.py", line 611, in two_qubit_decomposition
    num_cnots = _compute_num_cnots(U)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/ops/op_math/decompositions/two_qubit_unitary.py", line 126, in _compute_num_cnots
    evs = math.linalg.eigvals(gammaU)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autoray/autoray.py", line 79, in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/numpy/wrapper.py", line 117, in _wrapped
    res = obj(*args, **kwargs)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/tracer.py", line 45, in f_wrapped
    node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/core.py", line 34, in __init__
    raise NotImplementedError("VJP of {} wrt argnums {} not defined"
NotImplementedError: VJP of eigvals wrt argnums (0,) not defined

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/cvjjm/src/covqcstack/hfak/covvqetools/qubit_unitary_test.py", line 64, in <module>
    raise Exception("jac2 should not be nan") from exception
Exception: jac2 should not be nan

Compilation exited abnormally with code 1 at Tue Mar  5 11:51:09

System information

Name: PennyLane
Version: 0.34.0
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: /home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages
Editable project location: /home/cvjjm/src/covqcstack/hfak/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: covvqetools, PennyLane-Lightning

Platform info:           Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Python version:          3.9.16
Numpy version:           1.23.0
Scipy version:           1.10.0
Installed devices:
- cov.caching.qubit (covvqetools-0.7.0)
- cov.caching.qulacs (covvqetools-0.7.0)
- cov.parsec.casbox (covvqetools-0.7.0)
- cov.parsec.qubit (covvqetools-0.7.0)
- cov.qubit (covvqetools-0.7.0)
- cov.quicksilver.casbox (covvqetools-0.7.0)
- cov.quicksilver.qubit (covvqetools-0.7.0)
- default.gaussian (PennyLane-0.34.0)
- default.mixed (PennyLane-0.34.0)
- default.qubit (PennyLane-0.34.0)
- default.qubit.autograd (PennyLane-0.34.0)
- default.qubit.jax (PennyLane-0.34.0)
- default.qubit.legacy (PennyLane-0.34.0)
- default.qubit.tf (PennyLane-0.34.0)
- default.qubit.torch (PennyLane-0.34.0)
- default.qutrit (PennyLane-0.34.0)
- null.qubit (PennyLane-0.34.0)
- lightning.qubit (PennyLane-Lightning-0.34.0)

Existing GitHub issues

trbromley commented 7 months ago

Thanks @cvjjm! We're just finishing the last stages of our 0.35 release but will get back to this very soon.

albi3ro commented 7 months ago

It does look like there is source code notes for the two_qubit_decomposition:

https://github.com/PennyLaneAI/pennylane/blob/643df7b9cdd50baa3aeea6bec6b0ad0a401bb7f9/pennylane/ops/op_math/decompositions/two_qubit_unitary.py#L26

# I was not able to get this transform to be fully differentiable for unitary
# matrices that were constructed within a QNode, based on the QNode's input
# arguments. I would argue this is a fairly limited use case, but it would still
# be nice to have this eventually. Each interface fails for different reasons.
#
# - In Autograd, we obtain the AttributeError
#       'ArrayBox' object has no attribute 'conjugate'
#   for the 0-CNOT case when the zyz_decomposition function is called. In the
#   other cases, it cannot autodifferentiate through the linalg.eigvals function.

So this comes back to the repeating problem that we can't differentiate through eigendecomposition.

josh146 commented 7 months ago

Thanks for reporting this @cvjjm!

Finally, computing gradients also does not work. Differentiating at np.pi/2 yields a NotImplementedError in autograd (autograd==1.6.2). Differentiating at random parameters yields nan, even with diff_method="finite-diff".

The comment by @albi3ro above looks to be the cause --- @trbromley @albi3ro we should likely update the QubitUnitary docstring to include the decomposition gradient information if not already present.

  • For the given in_mat input matrix the matrix returned by qml.matrix() on a QNode that contains a single QubitUnitary operation is wrong and it is off by more than just a global phase. (For other input matrices such as, i.e., in_mat = make_unitary(np.pi/2 + 1e-6) the matrices match)
  • Also executing the QNode on default.qubit produces the wrong state, even though the operations on the tape look correct and when drawing the tape all looks good.

Interesting 🤔 Definitely something we need to look into.

cvjjm commented 7 months ago

Differentiation is not super high priority but getting the correct matrix under decomposition would be good :-)

But out of curiosity: Why does even finite-diff differentiation not work?

cvjjm commented 7 months ago

This is not the root cause of the issues described above, but it is in related code and may be worth improving as well:

https://github.com/PennyLaneAI/pennylane/blob/c12a29ca4009806272d45988328c14772f96280b/pennylane/ops/op_math/decompositions/single_qubit_unitary.py#L67-L78

The perturbation with epsilon seems to be not needed because:

>>> math.arctan2(0.1, 1e-64)
1.5707963267948966
>>> math.arctan2(0.1, 0)
1.5707963267948966
cvjjm commented 7 months ago

Hi again! I understand that there are a few more complex issues here that need more work and are maybe not that important, but it would be great if at least running circuits with QubitUnitary gates on default qubit would produce the correct states/expectation values.

isaacdevlugt commented 7 months ago

Hey @cvjjm! This has been on our radar, but thanks to @albi3ro we have a fix in the works 😄.

cvjjm commented 6 months ago

Does this only fix the (minor) problem of the discontinuity or do in_mat and out_mat in the above example now also match (at least up to a global phase)?

trbromley commented 6 months ago

Hi @cvjjm! Sorry for taking a while to get back to you. Running the above code with PL v0.36 coming out on Tuesday next week gives:

in_mat
[[ 1.   +0.j     0.   +0.j     0.   +0.j     0.   +0.j   ]
 [ 0.   +0.j     0.707-0.j     0.14 -0.693j  0.   +0.j   ]
 [ 0.   +0.j    -0.14 -0.693j  0.707+0.j     0.   +0.j   ]
 [ 0.   +0.j     0.   +0.j     0.   +0.j     1.   +0.j   ]]
out_mat
[[ 1.   +0.j     0.   -0.j    -0.   +0.j     0.   -0.j   ]
 [ 0.   -0.j     0.707-0.j     0.14 -0.693j -0.   +0.j   ]
 [-0.   +0.j    -0.14 -0.693j  0.707+0.j     0.   -0.j   ]
 [ 0.   -0.j     0.   +0.j    -0.   -0.j     1.   -0.j   ]]
global_phase 9.020562075079397e-16
out_mat with phase corrected
[[ 1.   +0.j     0.   -0.j    -0.   +0.j     0.   -0.j   ]
 [ 0.   -0.j     0.707-0.j     0.14 -0.693j -0.   +0.j   ]
 [-0.   +0.j    -0.14 -0.693j  0.707-0.j     0.   -0.j   ]
 [ 0.   -0.j     0.   +0.j    -0.   -0.j     1.   -0.j   ]]
difference should be zero but it is not:
[[ 0.+0.j -0.+0.j  0.-0.j -0.+0.j]
 [-0.+0.j -0.+0.j  0.+0.j  0.-0.j]
 [ 0.-0.j  0.+0.j -0.+0.j -0.+0.j]
 [-0.+0.j -0.-0.j  0.+0.j  0.+0.j]]
circuit:
0: ─╭U(M0)─┤ ╭State
1: ─╰U(M0)─┤ ╰State

M0 = 
[[ 1.        +0.00000000e+00j  0.        +0.00000000e+00j
   0.        +0.00000000e+00j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j  0.70710678-2.29934717e-17j
   0.14048043-6.93011723e-01j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j -0.14048043-6.93011723e-01j
   0.70710678+0.00000000e+00j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j  0.        +0.00000000e+00j
   0.        +0.00000000e+00j  1.        +0.00000000e+00j]]
tape looks correct:
[QubitUnitary(tensor([[ 1.        +0.00000000e+00j,  0.        +0.00000000e+00j,
          0.        +0.00000000e+00j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j,  0.70710678-2.29934717e-17j,
          0.14048043-6.93011723e-01j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j, -0.14048043-6.93011723e-01j,
          0.70710678+0.00000000e+00j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j,  0.        +0.00000000e+00j,
          0.        +0.00000000e+00j,  1.        +0.00000000e+00j]], requires_grad=True), wires=[0, 1])]
state returned by qnode should be |0> but it is not:
[ 1.00000000e+00+8.32667268e-16j  1.10929801e-16-1.10606113e-15j
 -7.85557000e-17+4.68148273e-16j  1.38777878e-17-1.38777878e-17j]
jac2 [nan]

This looks to be mostly working :rocket: However, we unfortunately are still getting nan for the Jacobian. I think the problem remains because of our inability to differentiate through the eigendecomposition as @albi3ro mentioned. Also, this is still a problem in finite-diff mode because that mode only applies finite difference to the quantum parts of the computation, the classical parts like eigendecomposition are still performed using automatic differentiation.