[BUG] Cannot pass a symbolic Keras tensor to postselect argument of qml.measure() when running on Tensorflow graph mode

[bear-2s]

Expected behavior

The variable can be accepted as the argument of the 'postselect' of qml.measure() function under the Tensorflow graph mode

Actual behavior

The variable cannot be accepted as the argument of the 'postselect' of qml.measure() function under the Tensorflow graph mode

Source code

import tensorflow as tf
import keras
from keras import mixed_precision
from keras.layers import Concatenate, concatenate
from keras.models import Model
import pennylane as qml
import numpy as np
from tensorflow.python.ops.numpy_ops import np_config
from tensorflow.experimental import numpy as tnp

tf.get_logger().setLevel('ERROR')
DEFAULT_TENSOR_TYPE = tf.float64
################################################################
#Parameters
################################################################
qbit = 2
bit = 2**qbit
num_epoch = 100000
LL = 40
QLL = 10
TT = 5
lamt = 0.1
sam = 1
################################################################
#Data preparation
################################################################
initial_data = tf.constant([[[ 0.70819672+0.j,         -0.43169441-0.14245724j,],
  [-0.43169441+0.14245724j,  0.29180328+0.j        ]]], dtype=tf.complex64)
initial_dataset = tf.data.Dataset.from_tensors(initial_data)
target_data = tf.constant([[[ 0.24980155+0.j,          0.40491117-0.15312634j],
  [ 0.40491117+0.15312634j,  0.75019845+0.j        ]]], dtype=tf.complex64)
target_dataset = tf.data.Dataset.from_tensors(target_data)
D_set = tf.data.Dataset.zip(initial_dataset,target_dataset)
################################################################
#Quantumm circuit
################################################################
dev = qml.device("default.qubit", wires=qbit+1)
@qml.qnode(dev, interface='tf')
def selecuit(inputs):
    qml.StatePrep(inputs, wires=0)
    return qml.probs(wires=[1])
@qml.qnode(dev, interface='tf')
def qcircuit(inputs,ps):
    qml.StatePrep(inputs, wires=0)
#Error occurs at the next line
    qml.measure(1,postselect = ps)
    return qml.density_matrix([0])
def convert_to_Svector(dm):
    a,b = tf.linalg.eigh(dm)
    c = tf.slice(b,[0,1],[2,1])
    c = tf.reshape(c,[2])
    return c
################################################################
#Custom Model
################################################################
class RNNModel(keras.Model):
    def __init__(self):
        super().__init__()
        self.mae_metric = keras.metrics.MeanAbsoluteError(name="mae")
    def train_step(self, data):
        qs, target = data
        loss = 0
        with tf.GradientTape() as tape:
            dm1 = self.data_flowing(qs[0,:,:],1)
        return { "mae": self.mae_metric.result()}
    def data_flowing(self,idm,nn):
        istate = convert_to_Svector(idm)
        istate = tf.convert_to_tensor(istate)
# selecuit is used to return the probability for the postselection result
        measure = selecuit(istate)
# measure is random picked 0 or 1 based on the probability
        measure = tf.reshape(measure, [1,2])
        measure = tf.random.categorical(tf.math.log(measure),1,dtype=tf.int32)
        measure = tf.constant([1], dtype=tf.int32)
        measure = tf.reshape(measure, [])
        odm = qcircuit(istate,measure)
        return odm
################################################################
#Excuting 
################################################################
model = RNNModel()
model.compile(optimizer="Adam")
ii = tf.ones([1,1,4], dtype=DEFAULT_TENSOR_TYPE)
state = tf.zeros([1,LL*qbit], dtype=DEFAULT_TENSOR_TYPE)
state1 = tf.zeros([1,LL*qbit+3], dtype=DEFAULT_TENSOR_TYPE)
model.fit(x = D_set, batch_size=None, epochs=num_epoch)

Tracebacks

Traceback (most recent call last):
  File "/home/max/qc/RNN/ps2.py", line 86, in <module>
    model.fit(x = D_set, batch_size=None, epochs=num_epoch)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_file1ew9_f6m.py", line 15, in tf__train_function
    retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
  File "/home/max/qc/RNN/ps2.py", line 64, in train_step
    dm1 = self.data_flowing(qs[0,:,:],1)
  File "/home/max/qc/RNN/ps2.py", line 76, in data_flowing
    odm = qcircuit(istate,measure)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/qnode.py", line 1164, in __call__
    return self._impl_call(*args, **kwargs)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/qnode.py", line 1150, in _impl_call
    res = self._execution_component(args, kwargs, override_shots=override_shots)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/qnode.py", line 1103, in _execution_component
    res = qml.execute(
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/execution.py", line 650, in execute
    tapes, post_processing = transform_program(tapes)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/transforms/core/transform_program.py", line 515, in __call__
    new_tapes, fn = transform(tape, *targs, **tkwargs)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/devices/preprocess.py", line 172, in mid_circuit_measurements
    return qml.defer_measurements(tape, device=device)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/transforms/core/transform_dispatcher.py", line 113, in __call__
    transformed_tapes, processing_fn = self._transform(obj, *targs, **tkwargs)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/transforms/defer_measurements.py", line 310, in defer_measurements
    new_operations.append(qml.Projector([op.postselect], wires=op.wires[0]))
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/capture/capture_meta.py", line 89, in __call__
    return type.__call__(cls, *args, **kwargs)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/ops/qubit/observables.py", line 452, in __init__
    state = tuple(qml.math.toarray(state).astype(int))
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/autoray/autoray.py", line 81, in do
    return func(*args, **kwargs)
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/autoray/autoray.py", line 1524, in numpy_to_numpy
    return do("asarray", x, like="numpy")
  File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/autoray/autoray.py", line 81, in do
    return func(*args, **kwargs)
NotImplementedError: in user code:

    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/keras/src/engine/training.py", line 1338, in train_function  *
        return step_function(self, iterator)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/keras/src/engine/training.py", line 1322, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/keras/src/engine/training.py", line 1303, in run_step  **
        outputs = model.train_step(data)
    File "/home/max/qc/RNN/ps2.py", line 64, in train_step
        dm1 = self.data_flowing(qs[0,:,:],1)
    File "/home/max/qc/RNN/ps2.py", line 76, in data_flowing
        odm = qcircuit(istate,measure)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/qnode.py", line 1164, in __call__
        return self._impl_call(*args, **kwargs)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/qnode.py", line 1150, in _impl_call
        res = self._execution_component(args, kwargs, override_shots=override_shots)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/qnode.py", line 1103, in _execution_component
        res = qml.execute(
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/workflow/execution.py", line 650, in execute
        tapes, post_processing = transform_program(tapes)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/transforms/core/transform_program.py", line 515, in __call__
        new_tapes, fn = transform(tape, *targs, **tkwargs)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/devices/preprocess.py", line 172, in mid_circuit_measurements
        return qml.defer_measurements(tape, device=device)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/transforms/core/transform_dispatcher.py", line 113, in __call__
        transformed_tapes, processing_fn = self._transform(obj, *targs, **tkwargs)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/transforms/defer_measurements.py", line 310, in defer_measurements
        new_operations.append(qml.Projector([op.postselect], wires=op.wires[0]))
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/capture/capture_meta.py", line 89, in __call__
        return type.__call__(cls, *args, **kwargs)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/pennylane/ops/qubit/observables.py", line 452, in __init__
        state = tuple(qml.math.toarray(state).astype(int))
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/autoray/autoray.py", line 81, in do
        return func(*args, **kwargs)
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/autoray/autoray.py", line 1524, in numpy_to_numpy
        return do("asarray", x, like="numpy")
    File "/root/miniconda3/envs/tc/lib/python3.9/site-packages/autoray/autoray.py", line 81, in do
        return func(*args, **kwargs)

    NotImplementedError: Cannot convert a symbolic tf.Tensor (Reshape_7:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported.

System information

Name: PennyLane
Version: 0.37.0
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: /root/miniconda3/envs/tc/lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Python version:          3.9.18
Numpy version:           1.24.3
Scipy version:           1.13.1
Installed devices:
- default.clifford (PennyLane-0.37.0)
- default.gaussian (PennyLane-0.37.0)
- default.mixed (PennyLane-0.37.0)
- default.qubit (PennyLane-0.37.0)
- default.qubit.autograd (PennyLane-0.37.0)
- default.qubit.jax (PennyLane-0.37.0)
- default.qubit.legacy (PennyLane-0.37.0)
- default.qubit.tf (PennyLane-0.37.0)
- default.qubit.torch (PennyLane-0.37.0)
- default.qutrit (PennyLane-0.37.0)
- default.qutrit.mixed (PennyLane-0.37.0)
- default.tensor (PennyLane-0.37.0)
- null.qubit (PennyLane-0.37.0)
- lightning.qubit (PennyLane-Lightning-0.37.0)
None

Existing GitHub issues

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

Related PennyLane Forum Questions:

[PL-5092]

[CatalinaAlbornoz]

Thanks for the bug report @bear-2s !

I’m linking the original Forum thread here for anyone else looking into this bug.