Fix tensorflow for when device returns float 32 results (#5446)

**Context:**

When testing the new lightning device, we discovered that TensorFlow
would error if the parameters were float64, the device returned float32,
and classical operations modified the parameter:

```
dev = qml.device('lightning.qubit', wires=2, c_dtype=np.complex64)

@qml.qnode(dev, diff_method="parameter-shift")
def circuit(x):
    qml.RX(tf.cos(x), wires=0)
    return qml.expval(qml.Z(0))

x = tf.Variable(0.1, dtype=tf.float64)

with tf.GradientTape() as tape:
    y = circuit(x)

tape.gradient(y, x)
```
```
InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a float tensor but is a double tensor [Op:Mul] name: 
```

The problem is that the output were results were `float32` precision, so
the vjp would be `float32`. But `float32` vjp cannot be combined with
`tf.cos`'s `float64` vjp.

The reverse problem (`float64` results but `float32` data) seems to be
fine.

**Description of the Change:**

If the input data are `float64` or `complex128`, we promote the results
to the type of the parameters.

[sc-58966]

---------

Co-authored-by: Mudit Pandey <[email protected]>

doc/releases/changelog-dev.md

@@ -303,6 +303,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Tensorflow can now handle devices with float32 results but float64 input parameters.
+  [(#5446)](https://github.com/PennyLaneAI/pennylane/pull/5446)
+
 * Fix a bug where the `argnum` kwarg of `qml.gradients.stoch_pulse_grad` references the wrong parameters in a tape,
   creating an inconsistency with other differentiation methods and preventing some use cases.
   [(#5458)](https://github.com/PennyLaneAI/pennylane/pull/5458)

pennylane/workflow/interfaces/tensorflow.py

@@ -118,7 +118,20 @@ def set_parameters_on_copy(tapes, params):
     return tuple(t.bind_new_parameters(a, list(range(len(a)))) for t, a in zip(tapes, params))
 
 
-def _to_tensors(x, dtype=None):
+def _get_parameters_dtype(parameters):
+    for p in parameters:
+        if qml.math.get_interface(p) == "tensorflow":
+            return p.dtype
+    return None
+
+
+_complex_dtype_map = {
+    tf.float32: tf.complex64,
+    tf.float64: tf.complex128,
+}
+
+
+def _to_tensors(x, dtype=None, complex_safe=False):
     """
     Convert a nested tuple structure of arrays into a nested tuple
     structure of TF tensors
@@ -128,8 +141,10 @@ def _to_tensors(x, dtype=None):
         return x
 
     if isinstance(x, (tuple, list)):
-        return tuple(_to_tensors(x_, dtype=dtype) for x_ in x)
+        return tuple(_to_tensors(x_, dtype=dtype, complex_safe=complex_safe) for x_ in x)
 
+    if complex_safe and "complex" in qml.math.get_dtype_name(x):
+        return tf.convert_to_tensor(x, dtype=_complex_dtype_map.get(dtype, dtype))
     return tf.convert_to_tensor(x, dtype=dtype)
 
 
@@ -212,7 +227,10 @@ def tf_execute(tapes, execute_fn, jpc, device=None, differentiable=False):
 
     # need to use same tapes for forward pass execution that we will use for the vjp
     # if we are using device derivatives (`not differentiable`) so we can find them in the cache
-    res = _to_tensors(execute_fn(numpy_tapes))
+    params_dtype = _get_parameters_dtype(parameters)
+    dtype = params_dtype if params_dtype in {tf.float64, tf.complex128} else None
+    # make sure is float64 if data is float64.  May cause errors otherwise if device returns float32 precision
+    res = _to_tensors(execute_fn(numpy_tapes), dtype=dtype, complex_safe=True)
 
     @tf.custom_gradient
     def custom_gradient_execute(*parameters):  # pylint:disable=unused-argument

tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py

@@ -811,3 +811,36 @@ def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, shots, u
         jac = qml.math.hstack(tape.jacobian(res, [weights, coeffs1, coeffs2]), like="tensorflow")
         expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
         assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0)
+
+
+@pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift"))
+def test_device_returns_float32(diff_method):
+    """Test that if the device returns float32, the derivative succeeds."""
+
+    def _to_float32(results):
+        if isinstance(results, (list, tuple)):
+            return tuple(_to_float32(r) for r in results)
+        return np.array(results, dtype=np.float32)
+
+    class Float32Dev(qml.devices.DefaultQubit):
+        def execute(self, circuits, execution_config=qml.devices.DefaultExecutionConfig):
+            results = super().execute(circuits, execution_config)
+            return _to_float32(results)
+
+    dev = Float32Dev()
+
+    @qml.qnode(dev, diff_method=diff_method)
+    def circuit(x):
+        qml.RX(tf.cos(x), wires=0)
+        return qml.expval(qml.Z(0))
+
+    x = tf.Variable(0.1, dtype=tf.float64)
+
+    with tf.GradientTape() as tape:
+        y = circuit(x)
+
+    assert qml.math.allclose(y, np.cos(np.cos(0.1)))
+
+    g = tape.gradient(y, x)
+    expected_g = np.sin(np.cos(0.1)) * np.sin(0.1)
+    assert qml.math.allclose(g, expected_g)