Repository: PennyLaneAI/pennylane

Issue:
[BUG] Taking gradients with `interface='jax'` or `interface='jax-python'` and backend=`qiskit.aer´ fails
### Expected behavior

The code should compute the gradients with jax.grad using the `qiskit.aer` backend. 

Interestingly, as @josh146 outlined, using @jit before the circuit or using the interface `jax-jit` instead of just `jax` works. This bug seems to be with the `jax-python` interface. You can change the interface to jax-jit or jax-python to check.

### Actual behavior

Error:
`    raise CircuitError(f"Invalid param type {type(parameter)} for gate {self.name}.")
qiskit.circuit.exceptions.CircuitError: "Invalid param type <class 'jaxlib.xla_extension.DeviceArray'> for gate rz."`

### Additional information

_No response_

### Source code

```shell
import pennylane as qml
from pennylane import numpy as np

import jax.numpy as jnp
from jax import grad
from jax import jit


N = 2
dev = qml.device("qiskit.aer", wires=N, shots=1000)

# WORKS with @jit or interface="jax-jit" but not with interface="jax-python"
@qml.qnode(dev, interface="jax")
def energy(params):
    qml.StronglyEntanglingLayers(params, wires=range(N))
    return qml.expval(qml.PauliZ(0))

params = np.random.random(qml.StronglyEntanglingLayers.shape(n_layers=2, n_wires=N))
params = jnp.array(params)
print("Energy:", energy(params))
print(grad(energy)(params))


print(qml.about())
```


### Tracebacks

```shell
Traceback (most recent call last):
  File "/Users/shahnawaz/Dropbox/phd/implicit-diff/quantum-implicit-differentiation/tests/test_jax_qiskit.py", line 21, in <module>
    print(grad(energy)(params))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/qnode.py", line 560, in __call__
    res = qml.execute(
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 412, in execute
    res = _execute(
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 110, in execute
    return _execute(
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 258, in _execute
    return wrapped_exec(params)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: qiskit.circuit.exceptions.CircuitError: "Invalid param type <class 'jaxlib.xla_extension.DeviceArray'> for gate rz."

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/shahnawaz/Dropbox/phd/implicit-diff/quantum-implicit-differentiation/tests/test_jax_qiskit.py", line 21, in <module>
    print(grad(energy)(params))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/api.py", line 991, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/api.py", line 1073, in value_and_grad_f
    g = vjp_py(jax.lax._one(ans))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/tree_util.py", line 279, in __call__
    return self.fun(*args, **kw)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/api.py", line 2428, in _vjp_pullback_wrapper
    ans = fun(*args)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/tree_util.py", line 279, in __call__
    return self.fun(*args, **kw)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/interpreters/ad.py", line 123, in unbound_vjp
    arg_cts = backward_pass(jaxpr, reduce_axes, consts, dummy_args, cts)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/interpreters/ad.py", line 229, in backward_pass
    cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/interpreters/ad.py", line 694, in _custom_lin_transpose
    cts_in = bwd.call_wrapped(*res, *cts_out)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/custom_derivatives.py", line 640, in <lambda>
    bwd_ = lu.wrap_init(lambda *args: bwd.call_wrapped(*args))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 204, in wrapped_exec_bwd
    partial_res = execute_fn(vjp_tapes)[0]
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 173, in wrapper
    res = fn(execution_tapes.values(), **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 125, in fn
    return original_fn(tapes, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 426, in batch_execute
    compiled_circuits = self.compile_circuits(circuits)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 415, in compile_circuits
    self.create_circuit_object(circuit.operations, rotations=circuit.diagonalizing_gates)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 228, in create_circuit_object
    applied_operations = self.apply_operations(operations)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 289, in apply_operations
    gate = mapped_operation(*par)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/library/standard_gates/rz.py", line 61, in __init__
    super().__init__("rz", 1, [phi], label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 40, in __init__
    super().__init__(name, num_qubits, 0, params, label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 100, in __init__
    self.params = params  # must be at last (other properties may be required for validation)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 216, in params
    self._params.append(self.validate_parameter(single_param))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 241, in validate_parameter
    raise CircuitError(f"Invalid param type {type(parameter)} for gate {self.name}.")
jax._src.traceback_util.UnfilteredStackTrace: qiskit.circuit.exceptions.CircuitError: "Invalid param type <class 'jaxlib.xla_extension.DeviceArray'> for gate rz."

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/shahnawaz/Dropbox/phd/implicit-diff/quantum-implicit-differentiation/tests/test_jax_qiskit.py", line 21, in <module>
    print(grad(energy)(params))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 204, in wrapped_exec_bwd
    partial_res = execute_fn(vjp_tapes)[0]
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 173, in wrapper
    res = fn(execution_tapes.values(), **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 125, in fn
    return original_fn(tapes, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 426, in batch_execute
    compiled_circuits = self.compile_circuits(circuits)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 415, in compile_circuits
    self.create_circuit_object(circuit.operations, rotations=circuit.diagonalizing_gates)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 228, in create_circuit_object
    applied_operations = self.apply_operations(operations)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 289, in apply_operations
    gate = mapped_operation(*par)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/library/standard_gates/rz.py", line 61, in __init__
    super().__init__("rz", 1, [phi], label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 40, in __init__
    super().__init__(name, num_qubits, 0, params, label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 100, in __init__
    self.params = params  # must be at last (other properties may be required for validation)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 216, in params
    self._params.append(self.validate_parameter(single_param))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 241, in validate_parameter
    raise CircuitError(f"Invalid param type {type(parameter)} for gate {self.name}.")
qiskit.circuit.exceptions.CircuitError: "Invalid param type <class 'jaxlib.xla_extension.DeviceArray'> for gate rz."
(mitiq) Shahnawazs-MacBook-Pro:tests shahnawaz$ python test_jax_qiskit.py 
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Energy: 0.242
Traceback (most recent call last):
  File "/Users/shahnawaz/Dropbox/phd/implicit-diff/quantum-implicit-differentiation/tests/test_jax_qiskit.py", line 21, in <module>
    print(grad(energy)(params))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/qnode.py", line 560, in __call__
    res = qml.execute(
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 412, in execute
    res = _execute(
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 110, in execute
    return _execute(
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 258, in _execute
    return wrapped_exec(params)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: qiskit.circuit.exceptions.CircuitError: "Invalid param type <class 'jaxlib.xla_extension.DeviceArray'> for gate rz."

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/shahnawaz/Dropbox/phd/implicit-diff/quantum-implicit-differentiation/tests/test_jax_qiskit.py", line 21, in <module>
    print(grad(energy)(params))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/api.py", line 991, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/api.py", line 1073, in value_and_grad_f
    g = vjp_py(jax.lax._one(ans))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/tree_util.py", line 279, in __call__
    return self.fun(*args, **kw)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/api.py", line 2428, in _vjp_pullback_wrapper
    ans = fun(*args)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/tree_util.py", line 279, in __call__
    return self.fun(*args, **kw)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/interpreters/ad.py", line 123, in unbound_vjp
    arg_cts = backward_pass(jaxpr, reduce_axes, consts, dummy_args, cts)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/interpreters/ad.py", line 229, in backward_pass
    cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/interpreters/ad.py", line 694, in _custom_lin_transpose
    cts_in = bwd.call_wrapped(*res, *cts_out)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/_src/custom_derivatives.py", line 640, in <lambda>
    bwd_ = lu.wrap_init(lambda *args: bwd.call_wrapped(*args))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 204, in wrapped_exec_bwd
    partial_res = execute_fn(vjp_tapes)[0]
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 173, in wrapper
    res = fn(execution_tapes.values(), **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 125, in fn
    return original_fn(tapes, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 426, in batch_execute
    compiled_circuits = self.compile_circuits(circuits)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 415, in compile_circuits
    self.create_circuit_object(circuit.operations, rotations=circuit.diagonalizing_gates)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 228, in create_circuit_object
    applied_operations = self.apply_operations(operations)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 289, in apply_operations
    gate = mapped_operation(*par)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/library/standard_gates/rz.py", line 61, in __init__
    super().__init__("rz", 1, [phi], label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 40, in __init__
    super().__init__(name, num_qubits, 0, params, label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 100, in __init__
    self.params = params  # must be at last (other properties may be required for validation)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 216, in params
    self._params.append(self.validate_parameter(single_param))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 241, in validate_parameter
    raise CircuitError(f"Invalid param type {type(parameter)} for gate {self.name}.")
jax._src.traceback_util.UnfilteredStackTrace: qiskit.circuit.exceptions.CircuitError: "Invalid param type <class 'jaxlib.xla_extension.DeviceArray'> for gate rz."

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/shahnawaz/Dropbox/phd/implicit-diff/quantum-implicit-differentiation/tests/test_jax_qiskit.py", line 21, in <module>
    print(grad(energy)(params))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/jax.py", line 204, in wrapped_exec_bwd
    partial_res = execute_fn(vjp_tapes)[0]
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 173, in wrapper
    res = fn(execution_tapes.values(), **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane/interfaces/batch/__init__.py", line 125, in fn
    return original_fn(tapes, **kwargs)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 426, in batch_execute
    compiled_circuits = self.compile_circuits(circuits)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 415, in compile_circuits
    self.create_circuit_object(circuit.operations, rotations=circuit.diagonalizing_gates)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 228, in create_circuit_object
    applied_operations = self.apply_operations(operations)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/pennylane_qiskit/qiskit_device.py", line 289, in apply_operations
    gate = mapped_operation(*par)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/library/standard_gates/rz.py", line 61, in __init__
    super().__init__("rz", 1, [phi], label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 40, in __init__
    super().__init__(name, num_qubits, 0, params, label=label)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 100, in __init__
    self.params = params  # must be at last (other properties may be required for validation)
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/instruction.py", line 216, in params
    self._params.append(self.validate_parameter(single_param))
  File "/Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages/qiskit/circuit/gate.py", line 241, in validate_parameter
    raise CircuitError(f"Invalid param type {type(parameter)} for gate {self.name}.")
qiskit.circuit.exceptions.CircuitError: "Invalid param type <class 'jaxlib.xla_extension.DeviceArray'> for gate rz."
```


### System information

```shell
WARNING: pip is being invoked by an old script wrapper. This will fail in a future version of pip.
Please see https://github.com/pypa/pip/issues/5599 for advice on fixing the underlying issue.
To avoid this problem you can invoke Python with '-m pip' instead of running pip directly.
Name: PennyLane
Version: 0.21.0
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: /Users/shahnawaz/miniconda3/envs/mitiq/lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning, PennyLane-qiskit
Platform info:           macOS-11.2.3-x86_64-i386-64bit
Python version:          3.9.10
Numpy version:           1.22.2
Scipy version:           1.7.3
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.21.0)
- qiskit.aer (PennyLane-qiskit-0.21.0)
- qiskit.basicaer (PennyLane-qiskit-0.21.0)
- qiskit.ibmq (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.21.0)
- default.gaussian (PennyLane-0.21.0)
- default.mixed (PennyLane-0.21.0)
- default.qubit (PennyLane-0.21.0)
- default.qubit.autograd (PennyLane-0.21.0)
- default.qubit.jax (PennyLane-0.21.0)
- default.qubit.tf (PennyLane-0.21.0)
- default.qubit.torch (PennyLane-0.21.0)
```


### Existing GitHub issues

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


Discussion:

Resolving patch:
diff --git a/pennylane/interfaces/batch/jax.py b/pennylane/interfaces/batch/jax.py
--- a/pennylane/interfaces/batch/jax.py
+++ b/pennylane/interfaces/batch/jax.py
@@ -193,15 +193,16 @@ def wrapped_exec_bwd(params, g):
                 new_tapes[-1].set_parameters(a)
                 new_tapes[-1].trainable_params = t.trainable_params
 
-            vjp_tapes, processing_fn = qml.gradients.batch_vjp(
-                new_tapes,
-                dy,
-                gradient_fn,
-                reduction="append",
-                gradient_kwargs=gradient_kwargs,
-            )
+            with qml.tape.Unwrap(*new_tapes):
+                vjp_tapes, processing_fn = qml.gradients.batch_vjp(
+                    new_tapes,
+                    dy,
+                    gradient_fn,
+                    reduction="append",
+                    gradient_kwargs=gradient_kwargs,
+                )
 
-            partial_res = execute_fn(vjp_tapes)[0]
+                partial_res = execute_fn(vjp_tapes)[0]
 
             for t in tapes:
                 multi_probs = (


Test patch:
diff --git a/tests/interfaces/test_batch_jax_qnode.py b/tests/interfaces/test_batch_jax_qnode.py
--- a/tests/interfaces/test_batch_jax_qnode.py
+++ b/tests/interfaces/test_batch_jax_qnode.py
@@ -409,12 +409,14 @@ def circuit(a, b):
         assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0)
         spy.assert_not_called()
 
-    def test_gradient_integration(self, interface, tol):
+    def test_gradient_integration(self, interface, tol, mocker):
         """Test that temporarily setting the shots works
         for gradient computations"""
         dev = qml.device("default.qubit", wires=2, shots=100)
         a, b = jnp.array([0.543, -0.654])
 
+        spy = mocker.spy(dev, "batch_execute")
+
         @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface)
         def cost_fn(a, b):
             qml.RY(a, wires=0)
@@ -427,6 +429,7 @@ def cost_fn(a, b):
 
         expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]
         assert np.allclose(res, expected, atol=0.1, rtol=0)
+        assert all(not isinstance(p, jnp.ndarray) for p in spy.call_args[0][0][0].get_parameters())
 
     def test_update_diff_method(self, mocker, interface, tol):
         """Test that temporarily setting the shots updates the diff method"""