Repository: PennyLaneAI/pennylane (python)

Issue:
Incorrect value of gradients at 0 after taking gradient twice
#### Issue description

A colleague noticed that when running `qml.grad` twice on a particular circuit, to compute a second derivative, the output value of the gradient is incorrect when the input value is precisely 0. The specific circuit is effectively the function `cos(x)` (see below for the code to reproduce).

* *Expected behavior:* The value of the second gradient at 0 for this function should be -1. 

* *Actual behavior:* The value of the second gradient at 0 is computed to be -0.5.

* *Reproduces how often:* Always 

* *System information:* (post the output of `import pennylane as qml; qml.about()`)

```
Name: PennyLane
Version: 0.15.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: /home/olivia/Software/anaconda3/envs/xanadu/lib/python3.8/site-packages
Requires: appdirs, networkx, scipy, autograd, toml, semantic-version, numpy
Required-by: PennyLane-SF, pennylane-qulacs, PennyLane-qsharp, PennyLane-qiskit, PennyLane-Qchem, PennyLane-Forest, PennyLane-Cirq
Platform info:           Linux-5.4.0-66-generic-x86_64-with-glibc2.10
Python version:          3.8.5
Numpy version:           1.18.5
Scipy version:           1.4.1
Installed devices:
- default.gaussian (PennyLane-0.15.0.dev0)
- default.mixed (PennyLane-0.15.0.dev0)
- default.qubit (PennyLane-0.15.0.dev0)
- default.qubit.autograd (PennyLane-0.15.0.dev0)
- default.qubit.jax (PennyLane-0.15.0.dev0)
- default.qubit.tf (PennyLane-0.15.0.dev0)
- default.tensor (PennyLane-0.15.0.dev0)
- default.tensor.tf (PennyLane-0.15.0.dev0)
- strawberryfields.fock (PennyLane-SF-0.13.0.dev0)
- strawberryfields.gaussian (PennyLane-SF-0.13.0.dev0)
- strawberryfields.gbs (PennyLane-SF-0.13.0.dev0)
- strawberryfields.remote (PennyLane-SF-0.13.0.dev0)
- strawberryfields.tf (PennyLane-SF-0.13.0.dev0)
- qulacs.simulator (pennylane-qulacs-0.12.0)
- microsoft.QuantumSimulator (PennyLane-qsharp-0.8.0)
- qiskit.aer (PennyLane-qiskit-0.14.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.14.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.14.0.dev0)
- forest.numpy_wavefunction (PennyLane-Forest-0.13.0.dev0)
- forest.qpu (PennyLane-Forest-0.13.0.dev0)
- forest.qvm (PennyLane-Forest-0.13.0.dev0)
- forest.wavefunction (PennyLane-Forest-0.13.0.dev0)
- cirq.mixedsimulator (PennyLane-Cirq-0.13.0)
- cirq.pasqal (PennyLane-Cirq-0.13.0)
- cirq.qsim (PennyLane-Cirq-0.13.0)
- cirq.qsimh (PennyLane-Cirq-0.13.0)
- cirq.simulator (PennyLane-Cirq-0.13.0)
```

#### Source code and tracebacks

Code to reproduce:

```
import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=1)

# Produces function cos(x)
@qml.qnode(dev)
def circuit(x):
    qml.RY(x, wires=0)
    return qml.expval(qml.PauliZ(0))

gradcos  = qml.grad(circuit)  # -sin(x)
gradcos2 = qml.grad(gradcos)  # -cos(x)

print(gradcos2(0.)) # Should be -1, but outputs -0.5

```

The issue seems to be at 0 specifically, as evaluating the gradient at `1e-6` yields a result very close to -1. 

I tried another example as well:
```
# Produces function -i*sin(x)
@qml.qnode(dev)
def circuit(x):
    qml.PauliX(wires=0)
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

gradcos  = qml.grad(circuit)  # -i * cos(x)
gradcos2 = qml.grad(gradcos)  # i * sin(x)

print(gradcos2(0.)) # Should be 1, but outputs 0.5
```


Pull request description:
Remove `abs` from `DefaultQubit`
**Context:**
This PR implements the fix to #1125 proposed in [this comment](https://github.com/PennyLaneAI/pennylane/issues/1125#issuecomment-999206195)

**Description of the Change:**
The line 792
```python
    prob = self.marginal_prob(self._abs(self._flatten(self._state)) ** 2, wires)
```
in `qml/devices/default_qubit.py` is replaced by
```python
    flat_state = self._flatten(self._state)
    real_state = self._real(flat_state)
    imag_state = self._imag(flat_state)
    prob = self.marginal_prob(real_state ** 2 + imag_state ** 2), wires)
```
avoiding the usage of `abs`, which does not have a well-defined derivative at 0.

**Benefits:**
Fixes #1125 in all interfaces.

**Possible Drawbacks:**
Performance, see short discussion [here](https://github.com/PennyLaneAI/pennylane/issues/1125#issuecomment-999206195)

**Related GitHub Issues:**
#1125


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 483da06aa..527bb68bd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -134,6 +134,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug in `DefaultQubit` where the second derivative of QNodes at 
+  positions corresponding to vanishing state vector amplitudes is wrong.
+  [(#2057)](https://github.com/PennyLaneAI/pennylane/pull/2057)
+
 * Fixes a bug where PennyLane didn't require v0.20.0 of PennyLane-Lightning,
   but raised an error with versions of Lightning earlier than v0.20.0 due to
   the new batch execution pipeline.
@@ -159,4 +163,4 @@
 
 This release contains contributions from (in alphabetical order):
 
-Juan Miguel Arrazola, Ali Asadi, Esther Cruz, Olivia Di Matteo, Diego Guala, Ankit Khandelwal, Antal Száva, David Wierichs, Shaoming Zhang
\ No newline at end of file
+Juan Miguel Arrazola, Ali Asadi, Esther Cruz, Olivia Di Matteo, Diego Guala, Ankit Khandelwal, Antal Száva, David Wierichs, Shaoming Zhang
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index a368407cc..1c8828d8c 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -789,5 +789,8 @@ class DefaultQubit(QubitDevice):
         if self._state is None:
             return None
 
-        prob = self.marginal_prob(self._abs(self._flatten(self._state)) ** 2, wires)
+        flat_state = self._flatten(self._state)
+        real_state = self._real(flat_state)
+        imag_state = self._imag(flat_state)
+        prob = self.marginal_prob(real_state ** 2 + imag_state ** 2, wires)
         return prob
diff --git a/pennylane/devices/default_qubit_autograd.py b/pennylane/devices/default_qubit_autograd.py
index 260ca58c8..0fa762e59 100644
--- a/pennylane/devices/default_qubit_autograd.py
+++ b/pennylane/devices/default_qubit_autograd.py
@@ -92,6 +92,7 @@ class DefaultQubitAutograd(DefaultQubit):
     _transpose = staticmethod(np.transpose)
     _tensordot = staticmethod(np.tensordot)
     _conj = staticmethod(np.conj)
+    _real = staticmethod(np.real)
     _imag = staticmethod(np.imag)
     _roll = staticmethod(np.roll)
     _stack = staticmethod(np.stack)
diff --git a/pennylane/devices/default_qubit_jax.py b/pennylane/devices/default_qubit_jax.py
index b54d59a83..1a218f077 100644
--- a/pennylane/devices/default_qubit_jax.py
+++ b/pennylane/devices/default_qubit_jax.py
@@ -148,6 +148,7 @@ class DefaultQubitJax(DefaultQubit):
         )
     )
     _conj = staticmethod(jnp.conj)
+    _real = staticmethod(jnp.real)
     _imag = staticmethod(jnp.imag)
     _roll = staticmethod(jnp.roll)
     _stack = staticmethod(jnp.stack)
diff --git a/pennylane/devices/default_qubit_tf.py b/pennylane/devices/default_qubit_tf.py
index 521384e47..d569dc7af 100644
--- a/pennylane/devices/default_qubit_tf.py
+++ b/pennylane/devices/default_qubit_tf.py
@@ -140,6 +140,7 @@ class DefaultQubitTF(DefaultQubit):
     _transpose = staticmethod(tf.transpose)
     _tensordot = staticmethod(tf.tensordot)
     _conj = staticmethod(tf.math.conj)
+    _real = staticmethod(tf.math.real)
     _imag = staticmethod(tf.math.imag)
     _roll = staticmethod(tf.roll)
     _stack = staticmethod(tf.stack)
diff --git a/pennylane/devices/default_qubit_torch.py b/pennylane/devices/default_qubit_torch.py
index 71838cbe0..30125d1c2 100644
--- a/pennylane/devices/default_qubit_torch.py
+++ b/pennylane/devices/default_qubit_torch.py
@@ -150,6 +150,7 @@ class DefaultQubitTorch(DefaultQubit):
     _transpose = staticmethod(lambda a, axes=None: a.permute(*axes))
     _asnumpy = staticmethod(lambda x: x.cpu().numpy())
     _conj = staticmethod(torch.conj)
+    _real = staticmethod(torch.real)
     _imag = staticmethod(torch.imag)
     _norm = staticmethod(torch.norm)
     _flatten = staticmethod(torch.flatten)


Test patch:
diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py
index 00d7e553a..a707d777c 100644
--- a/tests/devices/test_default_qubit_autograd.py
+++ b/tests/devices/test_default_qubit_autograd.py
@@ -280,6 +280,22 @@ class TestPassthruIntegration:
         )
         assert np.allclose(res, expected_grad, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("x, shift", [(0.0, 0.0), (0.5, -0.5)])
+    def test_hessian_at_zero(self, x, shift):
+        """Tests that the Hessian at vanishing state vector amplitudes
+        is correct."""
+        dev = qml.device("default.qubit.autograd", wires=1)
+
+        @qml.qnode(dev, interface="autograd", diff_method="backprop")
+        def circuit(x):
+            qml.RY(shift, wires=0)
+            qml.RY(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        assert qml.math.isclose(qml.jacobian(circuit)(x), 0.0)
+        assert qml.math.isclose(qml.jacobian(qml.jacobian(circuit))(x), -1.0)
+        assert qml.math.isclose(qml.grad(qml.grad(circuit))(x), -1.0)
+
     @pytest.mark.parametrize("operation", [qml.U3, qml.U3.decomposition])
     @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift", "finite-diff"])
     def test_autograd_interface_gradient(self, operation, diff_method, tol):
diff --git a/tests/devices/test_default_qubit_jax.py b/tests/devices/test_default_qubit_jax.py
index 9187771b7..4c2221268 100644
--- a/tests/devices/test_default_qubit_jax.py
+++ b/tests/devices/test_default_qubit_jax.py
@@ -543,6 +543,22 @@ class TestPassthruIntegration:
 
         assert jnp.allclose(jnp.array(res), jnp.array(expected_grad), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("x, shift", [(0.0, 0.0), (0.5, -0.5)])
+    def test_hessian_at_zero(self, x, shift):
+        """Tests that the Hessian at vanishing state vector amplitudes
+        is correct."""
+        dev = qml.device("default.qubit.jax", wires=1)
+
+        @qml.qnode(dev, interface="jax", diff_method="backprop")
+        def circuit(x):
+            qml.RY(shift, wires=0)
+            qml.RY(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        assert qml.math.isclose(jax.grad(circuit)(x), 0.0)
+        assert qml.math.isclose(jax.jacobian(jax.jacobian(circuit))(x), -1.0)
+        assert qml.math.isclose(jax.grad(jax.grad(circuit))(x), -1.0)
+
     @pytest.mark.parametrize("operation", [qml.U3, qml.U3.decomposition])
     @pytest.mark.parametrize("diff_method", ["backprop"])
     def test_jax_interface_gradient(self, operation, diff_method, tol):
diff --git a/tests/devices/test_default_qubit_tf.py b/tests/devices/test_default_qubit_tf.py
index 416060c43..dcb79f78d 100644
--- a/tests/devices/test_default_qubit_tf.py
+++ b/tests/devices/test_default_qubit_tf.py
@@ -1323,6 +1323,33 @@ class TestPassthruIntegration:
         res = tape.gradient(res, [a_tf, b_tf])
         assert np.allclose(res, expected_grad, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("x, shift", [(0.0, 0.0), (0.5, -0.5)])
+    def test_hessian_at_zero(self, x, shift):
+        """Tests that the Hessian at vanishing state vector amplitudes
+        is correct."""
+        dev = qml.device("default.qubit.tf", wires=1)
+
+        shift = tf.constant(shift)
+        x = tf.Variable(x)
+
+        @qml.qnode(dev, interface="tf", diff_method="backprop")
+        def circuit(x):
+            qml.RY(shift, wires=0)
+            qml.RY(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        with tf.GradientTape(persistent=True) as t2:
+            with tf.GradientTape(persistent=True) as t1:
+                value = circuit(x)
+            grad = t1.gradient(value, x)
+            jac = t1.jacobian(value, x)
+        hess_grad = t2.gradient(grad, x)
+        hess_jac = t2.jacobian(jac, x)
+
+        assert qml.math.isclose(grad, 0.0)
+        assert qml.math.isclose(hess_grad, -1.0)
+        assert qml.math.isclose(hess_jac, -1.0)
+
     @pytest.mark.parametrize("operation", [qml.U3, qml.U3.decomposition])
     @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift", "finite-diff"])
     def test_tf_interface_gradient(self, operation, diff_method, tol):
diff --git a/tests/devices/test_default_qubit_torch.py b/tests/devices/test_default_qubit_torch.py
index a3605f528..8e3371a2b 100644
--- a/tests/devices/test_default_qubit_torch.py
+++ b/tests/devices/test_default_qubit_torch.py
@@ -1425,6 +1425,26 @@ class TestPassthruIntegration:
         assert torch.allclose(a.grad, -0.5 * torch.sin(a) * (torch.cos(b) + 1), atol=tol, rtol=0)
         assert torch.allclose(b.grad, 0.5 * torch.sin(b) * (1 - torch.cos(a)))
 
+    @pytest.mark.parametrize("x, shift", [(0.0, 0.0), (0.5, -0.5)])
+    def test_hessian_at_zero(self, torch_device, x, shift):
+        """Tests that the Hessian at vanishing state vector amplitudes
+        is correct."""
+        dev = qml.device("default.qubit.torch", wires=1, torch_device=torch_device)
+
+        x = torch.tensor(x, requires_grad=True)
+
+        @qml.qnode(dev, interface="torch", diff_method="backprop")
+        def circuit(x):
+            qml.RY(shift, wires=0)
+            qml.RY(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        grad = torch.autograd.functional.jacobian(circuit, x)
+        hess = torch.autograd.functional.hessian(circuit, x)
+
+        assert qml.math.isclose(grad, torch.tensor(0.0))
+        assert qml.math.isclose(hess, torch.tensor(-1.0))
+
     @pytest.mark.parametrize("operation", [qml.U3, qml.U3.decomposition])
     @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift", "finite-diff"])
     def test_torch_interface_gradient(self, torch_device, operation, diff_method, tol):