[BUG] `qml.math.fidelity` is not differentiable for 2 or more wires

ludmilaasb commented 1 year ago

Expected behavior

The fidelity should be differentiable. For QNodes with a single wire, the gradient of fidelity is computed:

wires = 1
dev = qml.device('default.mixed', wires=wires)

@qml.qnode(dev)
def circuit1(x):
    qml.RX(x,wires=0)
    return qml.density_matrix(wires = [x for x in range(wires)])

@qml.qnode(dev)
def circuit2():
    return qml.density_matrix(wires = [x for x in range(wires)])

def cost(x):
    return qml.math.fidelity(circuit1(x),circuit2(),check_state=True)

grad_fn = qml.grad(cost)

x = np.tensor(0.5, requires_grad=True)
print("gradient: ",grad_fn(x))

gradient:  -0.6666666666666667

but returns some warnings, such as:

[/Users/ludmila.botelho/anaconda3/envs/qml/lib/python3.11/site-packages/autograd/numpy/numpy_vjps.py:99](https://file+.vscode-resource.vscode-cdn.net/Users/ludmila.botelho/anaconda3/envs/qml/lib/python3.11/site-packages/autograd/numpy/numpy_vjps.py:99): RuntimeWarning: divide by zero encountered in power
  defvjp(anp.sqrt,    lambda ans, x : lambda g: g * 0.5 * x**-0.5)
[/Users/ludmila.botelho/anaconda3/envs/qml/lib/python3.11/site-packages/autograd/numpy/numpy_wrapper.py:156](https://file+.vscode-resource.vscode-cdn.net/Users/ludmila.botelho/anaconda3/envs/qml/lib/python3.11/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)

Actual behavior

For 2 or more wires, the gradient fails and returns nan.

Additional information

This problem seems to be happening particularly for computing the gradient of fidelity between matrices with low rank, blocked, or sparse somehow. I also tried computing the gradient with jaxbut it didn't work:

grad_fn = jax.grad(cost)
print(grad_fn(0.3))

and the output:

Array(nan, dtype=float32, weak_type=True)

Source code

import pennylane as qml
import pennylane.numpy as np

wires = 2
dev = qml.device('default.mixed', wires=wires)

@qml.qnode(dev)
def circuit1(x):
    qml.RX(x,wires=0)
    return qml.density_matrix(wires = [x for x in range(wires)])

@qml.qnode(dev)
def circuit2():
    return qml.density_matrix(wires = [x for x in range(wires)])

def cost(x):
    return qml.math.fidelity(circuit1(x),circuit2(),check_state=True)

x = np.tensor(0.3, requires_grad=True)
grad_fn = qml.grad(cost)
print(grad_fn(x))

Tracebacks

~/anaconda3/envs/qml/lib/python3.11/site-packages/pennylane/math/quantum.py:46: UserWarning: Argument passed to fidelity has shape (4, 4) and will be interpreted as a density matrix. If a batched state vector was intended, please call qml.math.dm_from_state_vector first, as passing state vectors to fidelity is deprecated.
  warnings.warn(
~/anaconda3/envs/qml/lib/python3.11/site-packages/autograd/numpy/numpy_vjps.py:99: RuntimeWarning: invalid value encountered in power
  defvjp(anp.sqrt,    lambda ans, x : lambda g: g * 0.5 * x**-0.5)
~/anaconda3/envs/qml/lib/python3.11/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))
~anaconda3/envs/qml/lib/python3.11/site-packages/autograd/numpy/linalg.py:181: RuntimeWarning: invalid value encountered in multiply
  vjp_temp += _dot(_dot(vc, F * _dot(T(v), vg)), T(v))

System information

Name: PennyLane
Version: 0.31.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: ~/anaconda3/envs/qml/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning

Platform info:           macOS-13.4.1-arm64-arm-64bit
Python version:          3.11.3
Numpy version:           1.23.5
Scipy version:           1.10.0
Installed devices:
- default.gaussian (PennyLane-0.31.0)
- default.mixed (PennyLane-0.31.0)
- default.qubit (PennyLane-0.31.0)
- default.qubit.autograd (PennyLane-0.31.0)
- default.qubit.jax (PennyLane-0.31.0)
- default.qubit.tf (PennyLane-0.31.0)
- default.qubit.torch (PennyLane-0.31.0)
- default.qutrit (PennyLane-0.31.0)
- null.qubit (PennyLane-0.31.0)
- lightning.qubit (PennyLane-Lightning-0.31.0)

Existing GitHub issues

[X] I have searched existing GitHub issues to make sure the issue does not already exist.

ludmilaasb commented 1 year ago

One extra thing: unfortunately tests for the gradient of the fidelity between QNodes with wires $\geq 2$ are apparently lacking, I only spotted tests using just a single wire 😢

https://github.com/PennyLaneAI/pennylane/blob/9c19af26707ee77d467d2cb24f2531534923173d/tests/qinfo/test_fidelity.py#L198-L200

https://github.com/PennyLaneAI/pennylane/blob/9c19af26707ee77d467d2cb24f2531534923173d/tests/qinfo/test_fidelity.py#L246-L248

https://github.com/PennyLaneAI/pennylane/blob/9c19af26707ee77d467d2cb24f2531534923173d/tests/qinfo/test_fidelity.py#L265-L266

https://github.com/PennyLaneAI/pennylane/blob/9c19af26707ee77d467d2cb24f2531534923173d/tests/qinfo/test_fidelity.py#L690-L692

albi3ro commented 11 months ago

Is this issue now closed?

PennyLaneAI / pennylane