Repository: PennyLaneAI/pennylane
Pull request #2120 (2022-01-22 02:39:01)

Use `qml.math.cast` after `qml.math.convert_like` in `qml.gradients` to copy `dtype`
**Context:**
`qml.math.cast_like` attempts to copy the dtype of a tensor whereas `qml.math.convert_like` only uses the type of the tensors during conversion

The `qml.gradients` module uses `qml.math.convert_like`. There may, however, be cases where the dtype would need to be copied, otherwise unexpected behaviour arises.

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**
#2105 

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2120](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (ccafb79) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/6d165f9ee83f41c34048fd93318b1a625353d89b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (6d165f9) will **decrease** coverage by `2.12%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2120      +/-   ##
==========================================
- Coverage   99.19%   97.07%   -2.13%     
==========================================
  Files         228      228              
  Lines       17476    17485       +9     
==========================================
- Hits        17335    16973     -362     
- Misses        141      512     +371     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/param\_shift\_hessian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbV9zaGlmdF9oZXNzaWFuLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/vjp.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy92anAucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `84.41% <0.00%> (-15.59%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.73% <0.00%> (-12.27%)` | :arrow_down: |
| [pennylane/math/is\_independent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvaXNfaW5kZXBlbmRlbnQucHk=) | `88.50% <0.00%> (-11.50%)` | :arrow_down: |
| ... and [14 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [6d165f9...ccafb79](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2120?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6ce75d755..0b07c9790 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -426,6 +426,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug for the TensorFlow interface where the dtype of input tensors was
+  not cast.
+  [(#2120)](https://github.com/PennyLaneAI/pennylane/pull/2120)
+
 * Fixes a bug where batch transformed QNodes would fail to apply batch transforms
   provided by the underlying device.
   [(#2111)](https://github.com/PennyLaneAI/pennylane/pull/2111)
diff --git a/pennylane/gradients/param_shift_hessian.py b/pennylane/gradients/param_shift_hessian.py
index f3c8e39a5..e544fa0f0 100644
--- a/pennylane/gradients/param_shift_hessian.py
+++ b/pennylane/gradients/param_shift_hessian.py
@@ -60,7 +60,9 @@ def generate_multishifted_tapes(tape, idx, shifts):
         shifted_tape = tape.copy(copy_operations=True)
 
         for id_, s in zip(idx, shift):
+            dtype = new_params[id_].dtype
             new_params[id_] = new_params[id_] + qml.math.convert_like(s, new_params[id_])
+            new_params[id_] = qml.math.cast(new_params[id_], dtype)
 
         shifted_tape.set_parameters(new_params)
         tapes.append(shifted_tape)
@@ -171,9 +173,9 @@ def compute_hessian_tapes(tape, diff_methods, f0=None):
                 g = qml.math.zeros(out_dim)
             else:
                 res = qml.math.stack(res)
-                g = qml.math.tensordot(
-                    res, qml.math.convert_like(hessian_coeffs[k], res), [[0], [0]]
-                )
+                hess = qml.math.convert_like(hessian_coeffs[k], res)
+                hess = qml.math.cast(hess, res.dtype)
+                g = qml.math.tensordot(res, hess, [[0], [0]])
                 if (i, j) in unshifted_coeffs:
                     g += unshifted_coeffs[(i, j)] * r0
 
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
index 33334b14e..8164c1328 100644
--- a/pennylane/gradients/vjp.py
+++ b/pennylane/gradients/vjp.py
@@ -46,6 +46,7 @@ def compute_vjp(dy, jac, num=None):
 
     if not isinstance(dy_row, np.ndarray):
         jac = math.convert_like(jac, dy_row)
+        jac = math.cast(jac, dy_row.dtype)
 
     jac = math.reshape(jac, [num, -1])
 
@@ -54,7 +55,8 @@ def compute_vjp(dy, jac, num=None):
             # If the dy vector is zero, then the
             # corresponding element of the VJP will be zero.
             num_params = jac.shape[1]
-            return math.convert_like(np.zeros([num_params]), dy)
+            res = math.convert_like(np.zeros([num_params]), dy)
+            return math.cast(res, dy.dtype)
     except (AttributeError, TypeError):
         pass
 
@@ -169,7 +171,12 @@ def vjp(tape, dy, gradient_fn, gradient_kwargs=None):
             # If the dy vector is zero, then the
             # corresponding element of the VJP will be zero,
             # and we can avoid a quantum computation.
-            return [], lambda _, num=None: math.convert_like(np.zeros([num_params]), dy)
+
+            def func(_, num=None):  # pylint: disable=unused-argument
+                res = math.convert_like(np.zeros([num_params]), dy)
+                return math.cast(res, dy.dtype)
+
+            return [], func
     except (AttributeError, TypeError):
         pass
 


Test patch:
diff --git a/tests/gradients/test_vjp.py b/tests/gradients/test_vjp.py
index 001bb7f69..1fe4786f8 100644
--- a/tests/gradients/test_vjp.py
+++ b/tests/gradients/test_vjp.py
@@ -19,6 +19,7 @@ import pytest
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.gradients import param_shift
+import numpy as vanilla_numpy
 
 
 class TestComputeVJP:
@@ -51,6 +52,52 @@ class TestComputeVJP:
         vjp = qml.gradients.compute_vjp(dy, jac)
         assert np.all(vjp == np.zeros([3]))
 
+    @pytest.mark.parametrize("dtype1,dtype2", [("float32", "float64"), ("float64", "float32")])
+    def test_dtype_torch(self, dtype1, dtype2):
+        """Test that using the Torch interface the dtype of the result is
+        determined by the dtype of the dy."""
+        torch = pytest.importorskip("torch")
+
+        dtype1 = getattr(torch, dtype1)
+        dtype2 = getattr(torch, dtype2)
+
+        dy = torch.ones(4, dtype=dtype1)
+        jac = torch.ones((4, 4), dtype=dtype2)
+
+        assert qml.gradients.compute_vjp(dy, jac).dtype == dtype1
+
+    @pytest.mark.parametrize("dtype1,dtype2", [("float32", "float64"), ("float64", "float32")])
+    def test_dtype_tf(self, dtype1, dtype2):
+        """Test that using the TensorFlow interface the dtype of the result is
+        determined by the dtype of the dy."""
+        tf = pytest.importorskip("tensorflow")
+
+        dtype1 = getattr(tf, dtype1)
+        dtype2 = getattr(tf, dtype2)
+
+        dy = tf.ones(4, dtype=dtype1)
+        jac = tf.ones((4, 4), dtype=dtype2)
+
+        assert qml.gradients.compute_vjp(dy, jac).dtype == dtype1
+
+    @pytest.mark.parametrize("dtype1,dtype2", [("float32", "float64"), ("float64", "float32")])
+    def test_dtype_jax(self, dtype1, dtype2):
+        """Test that using the JAX interface the dtype of the result is
+        determined by the dtype of the dy."""
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+
+        dtype1 = getattr(jax.numpy, dtype1)
+        dtype2 = getattr(jax.numpy, dtype2)
+
+        dy = jax.numpy.array([0, 1], dtype=dtype1)
+        jac = jax.numpy.array([[0, 1], [2, 3]], dtype=dtype2)
+
+        assert qml.gradients.compute_vjp(dy, jac).dtype == dtype1
+
 
 class TestVJP:
     """Tests for the vjp function"""
@@ -187,6 +234,21 @@ class TestVJP:
 
         assert np.allclose(res, dy @ expected, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("dtype", [np.float32, np.float64])
+    def test_dtype_matches_dy(self, dtype):
+        """Tests that the vjp function matches the dtype of dy when dy is
+        zero-like."""
+        x = np.array([0.1], dtype=np.float64)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x[0], wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        dy = np.zeros(3, dtype=dtype)
+        _, func = qml.gradients.vjp(tape, dy, qml.gradients.param_shift)
+
+        assert func([]).dtype == dtype
+
 
 def expected(params):
     x, y = 1.0 * params
@@ -282,6 +344,40 @@ class TestVJPGradients:
         res = t.jacobian(vjp, params)
         assert np.allclose(res, qml.jacobian(expected)(params.numpy()), atol=tol, rtol=0)
 
+    def test_tf_custom_loss(self):
+        """Tests that the gradient pipeline using the TensorFlow interface with
+        a custom TF loss and lightning.qubit with a custom dtype does not raise
+        any errors."""
+        tf = pytest.importorskip("tensorflow")
+
+        nwires = 5
+        dev = qml.device("lightning.qubit", wires=nwires)
+        dev.C_DTYPE = vanilla_numpy.complex64
+        dev.R_DTYPE = vanilla_numpy.float32
+
+        @qml.qnode(dev, interface="tf", diff_method="adjoint")
+        def circuit(weights, features):
+            for i in range(nwires):
+                qml.RX(features[i], wires=i)
+                qml.RX(weights[i], wires=i)
+            return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]
+
+        vanilla_numpy.random.seed(42)
+
+        ndata = 100
+        data = [vanilla_numpy.random.randn(nwires).astype("float32") for _ in range(ndata)]
+        label = [vanilla_numpy.random.choice([1, 0]).astype("int") for _ in range(ndata)]
+
+        loss = tf.losses.SparseCategoricalCrossentropy()
+
+        params = tf.Variable(vanilla_numpy.random.randn(nwires).astype("float32"), trainable=True)
+        with tf.GradientTape() as tape:
+            probs = [circuit(params, d) for d in data]
+            loss_value = loss(label, probs)
+
+        grads = tape.gradient(loss_value, [params])
+        assert len(grads) == 1
+
     @pytest.mark.slow
     def test_jax(self, tol):
         """Tests that the output of the VJP transform