[BUG] tensorflow autograph with device derivatives and more than one measurement

[albi3ro]

Expected behavior

I would expect to be able to use tensorflow autograph to work with device derivatives.

Actual behavior

Various different error messages and failures.

Additional information

No response

Source code

@tf.function
@qml.qnode(
    qml.device('default.qubit', wires=2),
    diff_method="adjoint",
    grad_on_execution=False
)
def circuit(x):
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))

x = tf.Variable(0.5, dtype=tf.float64)

with tf.GradientTape() as tape:
    res = qml.math.hstack(circuit(x))
g = tape.jacobian(res, x)

Tracebacks

2024-01-17 12:23:05.885745: W tensorflow/core/framework/op_kernel.cc:1827] INVALID_ARGUMENT: TypeError: 'numpy.float64' object is not iterable
Traceback (most recent call last):

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/ops/script_ops.py", line 270, in __call__
    ret = func(*args)
          ^^^^^^^^^^^

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^

  File "/var/folders/k1/0v_kvphn55lgf_45kntf1hqm0000gq/T/__autograph_generated_file4mgs4va9.py", line 392, in _backward
    vjps = ag__.converted_call(ag__.ld(_compute_vjp), (ag__.ld(dy), ag__.ld(jac), ag__.ld(multi_measurements), ag__.ld(has_partitioned_shots)), None, fscope_7)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/autograph/impl/api.py", line 335, in converted_call
    return _call_unconverted(f, args, kwargs, options, False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/autograph/impl/api.py", line 460, in _call_unconverted
    return f(*args)
           ^^^^^^^^

  File "/Users/christina/Prog/pennylane/pennylane/workflow/interfaces/tensorflow.py", line 67, in _compute_vjp
    shot_vjps.append(qml.gradients.compute_vjp_multi(d, j))
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/christina/Prog/pennylane/pennylane/gradients/vjp.py", line 204, in compute_vjp_multi
    for d, j_ in zip(dy, jac):
                 ^^^^^^^^^^^^

TypeError: 'numpy.float64' object is not iterable


---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
Cell In[22], line 15
     13 with tf.GradientTape() as tape:
     14     res = qml.math.hstack(circuit(x))
---> 15 g = tape.jacobian(res, x)

File ~/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/eager/backprop.py:1169, in GradientTape.jacobian(self, target, sources, unconnected_gradients, parallel_iterations, experimental_use_pfor)
   1167 if experimental_use_pfor:
   1168   try:
-> 1169     output = pfor_ops.pfor(loop_fn, target_size,
   1170                            parallel_iterations=parallel_iterations)
   1171   except ValueError as err:
   1172     raise ValueError(
   1173         "Encountered an exception while vectorizing the "
   1174         "jacobian computation. Vectorization can be disabled by setting"
   1175         " experimental_use_pfor to False.") from err

File ~/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/ops/parallel_for/control_flow_ops.py:228, in pfor(loop_fn, iters, fallback_to_while_loop, parallel_iterations, warn)
    225     def_function.run_functions_eagerly(False)
    226   f = def_function.function(f)
--> 228 outputs = f()
    229 if functions_run_eagerly is not None:
    230   def_function.run_functions_eagerly(functions_run_eagerly)

File ~/Prog/pl311/lib/python3.11/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 ~/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/eager/execute.py:60, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     53   # Convert any objects of type core_types.Tensor to Tensor.
     54   inputs = [
     55       tensor_conversion_registry.convert(t)
     56       if isinstance(t, core_types.Tensor)
     57       else t
     58       for t in inputs
     59   ]
---> 60   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     61                                       inputs, attrs, num_outputs)
     62 except core._NotOkStatusException as e:
     63   if name is not None:

InvalidArgumentError: Graph execution error:

Detected at node gradients/IdentityN_grad/PyFunc/pfor/while/PyFunc defined at (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main

  File "<frozen runpy>", line 88, in _run_code

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel_launcher.py", line 17, in <module>

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/traitlets/config/application.py", line 1053, in launch_instance

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel/kernelapp.py", line 737, in start

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tornado/platform/asyncio.py", line 195, in start

  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/base_events.py", line 604, in run_forever

  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/base_events.py", line 1909, in _run_once

  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/events.py", line 80, in _run

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel/kernelbase.py", line 524, in dispatch_queue

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel/kernelbase.py", line 513, in process_one

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel/kernelbase.py", line 418, in dispatch_shell

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel/kernelbase.py", line 758, in execute_request

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel/ipkernel.py", line 426, in do_execute

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/ipykernel/zmqshell.py", line 549, in run_cell

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3024, in run_cell

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3079, in _run_cell

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3284, in run_cell_async

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3466, in run_ast_nodes

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3526, in run_code

  File "/var/folders/k1/0v_kvphn55lgf_45kntf1hqm0000gq/T/ipykernel_38901/3294547231.py", line 15, in <module>

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/ops/parallel_for/control_flow_ops.py", line 212, in f

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/ops/parallel_for/pfor.py", line 4526, in f

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/ops/parallel_for/pfor.py", line 4526, in f

TypeError: 'numpy.float64' object is not iterable
Traceback (most recent call last):

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/ops/script_ops.py", line 270, in __call__
    ret = func(*args)
          ^^^^^^^^^^^

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^

  File "/var/folders/k1/0v_kvphn55lgf_45kntf1hqm0000gq/T/__autograph_generated_file4mgs4va9.py", line 392, in _backward
    vjps = ag__.converted_call(ag__.ld(_compute_vjp), (ag__.ld(dy), ag__.ld(jac), ag__.ld(multi_measurements), ag__.ld(has_partitioned_shots)), None, fscope_7)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/autograph/impl/api.py", line 335, in converted_call
    return _call_unconverted(f, args, kwargs, options, False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/christina/Prog/pl311/lib/python3.11/site-packages/tensorflow/python/autograph/impl/api.py", line 460, in _call_unconverted
    return f(*args)
           ^^^^^^^^

  File "/Users/christina/Prog/pennylane/pennylane/workflow/interfaces/tensorflow.py", line 67, in _compute_vjp
    shot_vjps.append(qml.gradients.compute_vjp_multi(d, j))
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/christina/Prog/pennylane/pennylane/gradients/vjp.py", line 204, in compute_vjp_multi
    for d, j_ in zip(dy, jac):
                 ^^^^^^^^^^^^

TypeError: 'numpy.float64' object is not iterable


	 [[{{node gradients/IdentityN_grad/PyFunc/pfor/while/PyFunc}}]] [Op:__inference_f_2501]

System information

PL master

Existing GitHub issues

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