Hi Shawn,
The trace of the system, \text{Tr}(\rho), can also be written as \text{Tr}(\rho I)=\langle I\rangle, so we can equivalently think of it as the expectation of the identity operator.
This allows you to construct a QNode that returns the trace like so:
@qml.qnode(dev)
def layer(inputs, w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10):
qml.templates.SqueezingEmbedding(inputs, wires=range(wires))
qml.templates.CVNeuralNetLayers(w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10,wires=range(wires))
return qml.Identity(wires=range(wires))
Alternatively, without modifying your existing QNode, you can inspect the device after QNode evaluation to find the trace:
@qml.qnode(dev)
def layer(inputs, w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10):
qml.templates.SqueezingEmbedding(inputs, wires=range(wires))
qml.templates.CVNeuralNetLayers(w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10,wires=range(wires))
return [qml.expval(qml.X(wires=i)) for i in range(wires)]
# evaluate the QNode
result = layers(**inputs)
# Check the device trace
dev.state.trace()