[BUG] Tensorflow Dtype is not preserved with lightning.qubit

[mudit2812]

Expected behavior

Dtype is preserved and no errors are raised.

Actual behavior

Dtype is not preserved, leading to errors as TF is extremely strict with dtypes.

Additional information

Found the following failing test in the device test suite:

Source code

import pennylane as qml
import tensorflow as tf
import numpy as np
tol = 1e-7

wires = 3
dev = qml.device("lightning.qubit", wires=wires)

@qml.qnode(dev, diff_method="hadamard", max_diff=2)
def circuit(x):
    qml.RY(x[0], wires=0)
    qml.RX(x[1], wires=0)
    return qml.expval(qml.Z(0))

x = tf.Variable([1.0, 2.0], dtype=tf.float64)

with tf.GradientTape() as tape1:
    with tf.GradientTape() as tape2:
        res = circuit(x)
    g = tape2.gradient(res, x)

hess = tape1.jacobian(g, x)
a, b = x

expected_res = np.cos(a) * np.cos(b)
assert np.allclose(res, expected_res, atol=tol, rtol=0)

expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)]
assert np.allclose(g, expected_g, atol=tol, rtol=0)

expected_hess = [
    [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)],
    [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)],
]
assert np.allclose(hess, expected_hess, atol=tol, rtol=0)

Tracebacks

{
	"name": "ValueError",
	"message": "Tensor conversion requested dtype float64 for Tensor with dtype float32: <tf.Tensor: shape=(), dtype=float32, numpy=-0.2248451>",
	"stack": "---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 19
     17 with tf.GradientTape() as tape1:
     18     with tf.GradientTape() as tape2:
---> 19         res = circuit(x)
     20         print(res.dtype)
     21     g = tape2.gradient(res, x)

File ~/repos/pennylane/pennylane/workflow/qnode.py:1085, in QNode.__call__(self, *args, **kwargs)
   1082 self._update_gradient_fn(shots=override_shots, tape=self._tape)
   1084 try:
-> 1085     res = self._execution_component(args, kwargs, override_shots=override_shots)
   1086 finally:
   1087     if old_interface == \"auto\":

File ~/repos/pennylane/pennylane/workflow/qnode.py:1039, in QNode._execution_component(self, args, kwargs, override_shots)
   1036 full_transform_program.prune_dynamic_transform()
   1038 # pylint: disable=unexpected-keyword-arg
-> 1039 res = qml.execute(
   1040     (self._tape,),
   1041     device=self.device,
   1042     gradient_fn=self.gradient_fn,
   1043     interface=self.interface,
   1044     transform_program=full_transform_program,
   1045     config=config,
   1046     gradient_kwargs=self.gradient_kwargs,
   1047     override_shots=override_shots,
   1048     **self.execute_kwargs,
   1049 )
   1050 res = res[0]
   1052 # convert result to the interface in case the qfunc has no parameters

File ~/repos/pennylane/pennylane/workflow/execution.py:792, in execute(tapes, device, gradient_fn, interface, transform_program, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp)
    784 ml_boundary_execute = _get_ml_boundary_execute(
    785     interface,
    786     _grad_on_execution,
    787     config.use_device_jacobian_product,
    788     differentiable=max_diff > 1,
    789 )
    791 if interface in jpc_interfaces:
--> 792     results = ml_boundary_execute(tapes, execute_fn, jpc, device=device)
    793 else:
    794     results = ml_boundary_execute(
    795         tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff
    796     )

File ~/repos/pennylane/pennylane/workflow/interfaces/tensorflow.py:235, in tf_execute(tapes, execute_fn, jpc, device, differentiable)
    233 # make sure is float64 if data is float64.  May cause errors otherwise if device returns float32 precision
    234 print(dtype)
--> 235 res = _to_tensors(execute_fn(numpy_tapes), dtype=dtype, complex_safe=True)
    237 @tf.custom_gradient
    238 def custom_gradient_execute(*parameters):  # pylint:disable=unused-argument
    239     \"\"\"An execution of tapes with VJP's registered with tensorflow.
    240 
    241     Args:
   (...)
    252 
    253     \"\"\"

File ~/repos/pennylane/pennylane/workflow/interfaces/tensorflow.py:144, in _to_tensors(x, dtype, complex_safe)
    141     return x
    143 if isinstance(x, (tuple, list)):
--> 144     return tuple(_to_tensors(x_, dtype=dtype, complex_safe=complex_safe) for x_ in x)
    146 if complex_safe and \"complex\" in qml.math.get_dtype_name(x):
    147     return tf.convert_to_tensor(x, dtype=_complex_dtype_map.get(dtype, dtype))

File ~/repos/pennylane/pennylane/workflow/interfaces/tensorflow.py:144, in <genexpr>(.0)
    141     return x
    143 if isinstance(x, (tuple, list)):
--> 144     return tuple(_to_tensors(x_, dtype=dtype, complex_safe=complex_safe) for x_ in x)
    146 if complex_safe and \"complex\" in qml.math.get_dtype_name(x):
    147     return tf.convert_to_tensor(x, dtype=_complex_dtype_map.get(dtype, dtype))

File ~/repos/pennylane/pennylane/workflow/interfaces/tensorflow.py:148, in _to_tensors(x, dtype, complex_safe)
    146 if complex_safe and \"complex\" in qml.math.get_dtype_name(x):
    147     return tf.convert_to_tensor(x, dtype=_complex_dtype_map.get(dtype, dtype))
--> 148 return tf.convert_to_tensor(x, dtype=dtype)

File ~/.pyenv/versions/3.10.12/envs/lightning/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File ~/.pyenv/versions/3.10.12/envs/lightning/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:1599, in convert_to_tensor(value, dtype, name, as_ref, preferred_dtype, dtype_hint, ctx, accepted_result_types)
   1597 if isinstance(value, Tensor):
   1598   if dtype is not None and not dtype.is_compatible_with(value.dtype):
-> 1599     raise ValueError(
   1600         _add_error_prefix(
   1601             f\"Tensor conversion requested dtype {dtype.name} \"
   1602             f\"for Tensor with dtype {value.dtype.name}: {value!r}\",
   1603             name=name))
   1604   return value
   1606 if preferred_dtype is not None:

ValueError: Tensor conversion requested dtype float64 for Tensor with dtype float32: <tf.Tensor: shape=(), dtype=float32, numpy=-0.2248451>"
}

System information

Name: PennyLane
Version: 0.36.0.dev0
Summary: 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.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/mudit.pandey/repos/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing_extensions
Required-by: PennyLane_Lightning

Platform info:           macOS-14.4.1-arm64-arm-64bit
Python version:          3.10.12
Numpy version:           1.26.4
Scipy version:           1.13.0
Installed devices:
- default.clifford (PennyLane-0.36.0.dev0)
- default.gaussian (PennyLane-0.36.0.dev0)
- default.mixed (PennyLane-0.36.0.dev0)
- default.qubit (PennyLane-0.36.0.dev0)
- default.qubit.autograd (PennyLane-0.36.0.dev0)
- default.qubit.jax (PennyLane-0.36.0.dev0)
- default.qubit.legacy (PennyLane-0.36.0.dev0)
- default.qubit.tf (PennyLane-0.36.0.dev0)
- default.qubit.torch (PennyLane-0.36.0.dev0)
- default.qutrit (PennyLane-0.36.0.dev0)
- null.qubit (PennyLane-0.36.0.dev0)
- lightning.qubit (PennyLane_Lightning-0.36.0.dev21)

Existing GitHub issues

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