Using the state vector directly

Hi @kareem_essafty, the master version of PennyLane on GitHub has a new feature called the PassthruQNode. Using this QNode, with the default.tensor.tf device, should do what you want. For example, consider the following:

import tensorflow as tf
import pennylane as qml
from pennylane.qnodes import PassthruQNode

dev = qml.device('default.tensor.tf', wires=2)


def circuit(params):
    qml.RX(params[0], wires=0)
    qml.RX(params[1], wires=1)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(0))

qnode = PassthruQNode(circuit, dev)
params = tf.Variable([0.3, 0.1])


with tf.GradientTape() as tape:
    tape.watch(params)
    qnode(params)
    state = dev._state

grad = tape.gradient(state, params)

print("State:", state)
print("Gradient:", grad)

This gives the output:

State: tf.Tensor(
[[ 0.98753537+0.j          0.        -0.04941796j]
 [-0.00746879+0.j          0.        -0.14925138j]], shape=(2, 2), dtype=complex128)
Gradient: tf.Tensor([-0.09933467 -0.09933467], shape=(2,), dtype=float32)
2 Likes