[BUG] Error for user-defined device `dtype` with the TF interface

antalszava commented 2 years ago

Expected behavior

No error arising related to dtype mismatch when setting custom dtypes on the device and using TF.

Actual behavior

The following error is raised with a QNode that takes floats as input and a TF loss:

InvalidArgumentError: cannot compute AddN as input #10(zero-based) was expected to be a float tensor but is a double tensor [Op:AddN]

Additional information

There might be a conversion happening in the backward pass (potentially in the custom gradient implementation). When omitting the use of SparseCategoricalCrossentropy, the gradient can be obtained without an issue.

Source code

import tensorflow as tf
import numpy as np
import pennylane as qml

nwires = 5
dev = qml.device("lightning.qubit", wires=nwires)
dev.C_DTYPE = np.complex64
dev.R_DTYPE = np.float32

@qml.qnode(dev, interface='tf', diff_method='adjoint')
def circuit(weights, features=np.zeros(nwires)):
    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))]

np.random.seed(42)

ndata = 100
data = [np.random.randn(nwires).astype('float32') for _ in range(ndata)]
label = [np.random.choice([1, 0]).astype('int') for _ in range(ndata)]

loss = tf.losses.SparseCategoricalCrossentropy()

params = tf.Variable(np.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])

Tracebacks

No response

System information

Name: PennyLane
Version: 0.21.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/antal/xanadu/pennylane
Requires: numpy, scipy, networkx, retworkx, autograd, toml, appdirs, semantic_version, autoray, cachetools, pennylane-lightning
Required-by: PennyLane-Lightning, PennyLane-Cirq, PennyLane-Orquestra, pennylane-qulacs, amazon-braket-pennylane-plugin, PennyLane-Honeywell, PennyLane-qiskit, PennyLane-AQT, PennyLane-PQ, PennyLane-Forest, PennyLane-qsharp, PennyLane-Qchem, PennyLane-IonQ, PennyLane-SF
Platform info:           Linux-5.11.0-44-generic-x86_64-with-glibc2.10
Python version:          3.8.5
Numpy version:           1.21.4
Scipy version:           1.7.3
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.20.1)
- cirq.mixedsimulator (PennyLane-Cirq-0.19.0)
- cirq.pasqal (PennyLane-Cirq-0.19.0)
- cirq.qsim (PennyLane-Cirq-0.19.0)
- cirq.qsimh (PennyLane-Cirq-0.19.0)
- cirq.simulator (PennyLane-Cirq-0.19.0)
- orquestra.forest (PennyLane-Orquestra-0.15.0)
- orquestra.ibmq (PennyLane-Orquestra-0.15.0)
- orquestra.qiskit (PennyLane-Orquestra-0.15.0)
- orquestra.qulacs (PennyLane-Orquestra-0.15.0)
- qulacs.simulator (pennylane-qulacs-0.17.0.dev0)
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- honeywell.hqs (PennyLane-Honeywell-0.16.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.18.0.dev0)
- aqt.noisy_sim (PennyLane-AQT-0.18.0)
- aqt.sim (PennyLane-AQT-0.18.0)
- projectq.classical (PennyLane-PQ-0.18.0.dev0)
- projectq.ibm (PennyLane-PQ-0.18.0.dev0)
- projectq.simulator (PennyLane-PQ-0.18.0.dev0)
- forest.numpy_wavefunction (PennyLane-Forest-0.18.0.dev0)
- forest.qvm (PennyLane-Forest-0.18.0.dev0)
- forest.wavefunction (PennyLane-Forest-0.18.0.dev0)
- microsoft.QuantumSimulator (PennyLane-qsharp-0.19.0)
- ionq.qpu (PennyLane-IonQ-0.17.0.dev0)
- ionq.simulator (PennyLane-IonQ-0.17.0.dev0)
- strawberryfields.fock (PennyLane-SF-0.20.0.dev0)
- strawberryfields.gaussian (PennyLane-SF-0.20.0.dev0)
- strawberryfields.gbs (PennyLane-SF-0.20.0.dev0)
- strawberryfields.remote (PennyLane-SF-0.20.0.dev0)
- strawberryfields.tf (PennyLane-SF-0.20.0.dev0)
- default.gaussian (PennyLane-0.21.0.dev0)
- default.mixed (PennyLane-0.21.0.dev0)
- default.qubit (PennyLane-0.21.0.dev0)
- default.qubit.autograd (PennyLane-0.21.0.dev0)
- default.qubit.jax (PennyLane-0.21.0.dev0)
- default.qubit.tf (PennyLane-0.21.0.dev0)
- default.qubit.torch (PennyLane-0.21.0.dev0)

Existing GitHub issues

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

yitchen-tim commented 2 years ago

Adding info about the bug: It only happens with big number of ndata. For example, there is no bug when ndata=10 instead of 100. The value of ndata which starts triggering the bug also depends on the random seed. When the seed is 42, the bug happens when ndata is bigger than 14 on my machine.

josh146 commented 2 years ago

Thanks @yitchen-tim for the details! Although I find it very odd that it is dependent on the size of the features 🤔

@antalszava is it like the case that there are more float64 assumptions hardcoded in somewhere?

PennyLaneAI / pennylane