Repository: PennyLaneAI/pennylane
Pull request #1491 (2021-08-03 15:57:10)

Add a tape unwrapping context manager
**Context:** In #1490, improved handling for tensor unwrapping and determining trainability was added to PennyLane. A big application for this functionality is in the PL interfaces; in particular, the interfaces typically must perform the following logic:

- Extract the tape parameters
- Determine which are trainable
- Unwrap all parameters to NumPy arrays and floats
- Set the tape parameters as unwrapped
- Execute the device
- Set the original tape parameters

This PR provides a context manager to the tape to automate this procedure for any autodiff framework.

**Description of the Change:**

A new tape method, `tape.unwrap()` is added. This method is a context manager; inside the context, the tapes parameters are unwrapped to NumPy arrays and floats, the the trainable parameter indices are set.

These changes are temporary, and reverted on exiting the context.

Example:

```pycon
>>> with tf.GradientTape():
...     with qml.tape.QuantumTape() as tape:
...         qml.RX(tf.Variable(0.1), wires=0)
...         qml.RY(tf.constant(0.2), wires=0)
...         qml.RZ(tf.Variable(0.3), wires=0)
...     with UnwrapTape(tape) as unwrapped_tape:
...         print("Trainable params:", unwrapped_tape.trainable_params)
...         print("Unwrapped params:", unwrapped_tape.get_parameters())
Trainable params: {0, 2}
Unwrapped params: [0.1, 0.3]
>>> print("Original parameters:", tape.get_parameters())
Original parameters: [<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.1>,
  <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.3>]
```

**Benefits:**

- This logic, which was previously duplicated across all interfaces, is now provided as a separate, well tested, and easily extendible unit.

- Due to the previous code duplication, bugs fixed in some interfaces were not always found and fixed in others.

- The logic is now well tested with unit tests

- It is possible to use this unwrapping capability elsewhere, outside the interfaces. For example, some plugins may need to unwrap tape parameters in certain cases, see https://github.com/PennyLaneAI/pennylane/pull/893.

**Possible Drawbacks:** None, as far as I can tell.

**Related GitHub Issues:** n/a


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1491](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (424d19e) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/f2bb1a0d8941f68b59ba13ff77e32ad3de455dff?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (f2bb1a0) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1491   +/-   ##
=======================================
  Coverage   98.32%   98.33%           
=======================================
  Files         180      181    +1     
  Lines       12741    12762   +21     
=======================================
+ Hits        12528    12549   +21     
  Misses        213      213           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/math/single\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvc2luZ2xlX2Rpc3BhdGNoLnB5) | `99.46% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/tape/tape.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvdGFwZS5weQ==) | `97.80% <100.00%> (+0.01%)` | :arrow_up: |
| [pennylane/tape/unwrap.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvdW53cmFwLnB5) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [f2bb1a0...424d19e](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1491?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: [ch7187]

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index bb99020d8..27bdd37fc 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -423,6 +423,35 @@
 
   In addition, the behaviour of `qml.math.requires_grad` has been improved in order to
   correctly determine trainability during Autograd and JAX backwards passes.
+
+* A new tape method, `tape.unwrap()` is added. This method is a context manager; inside the
+  context, the tapes parameters are unwrapped to NumPy arrays and floats, and the trainable
+  parameter indices are set.
+  [(#1491)](https://github.com/PennyLaneAI/pennylane/pull/1491)
+
+  These changes are temporary, and reverted on exiting the context.
+
+  ```pycon
+  >>> with tf.GradientTape():
+  ...     with qml.tape.QuantumTape() as tape:
+  ...         qml.RX(tf.Variable(0.1), wires=0)
+  ...         qml.RY(tf.constant(0.2), wires=0)
+  ...         qml.RZ(tf.Variable(0.3), wires=0)
+  ...     with tape.unwrap():
+  ...         print("Trainable params:", tape.trainable_params)
+  ...         print("Unwrapped params:", tape.get_parameters())
+  Trainable params: {0, 2}
+  Unwrapped params: [0.1, 0.3]
+  >>> print("Original parameters:", tape.get_parameters())
+  Original parameters: [<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.1>,
+    <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.3>]
+  ```
+
+  In addition, ``qml.tape.Unwrap`` is a context manager that unwraps multiple tapes:
+
+  ```pycon
+  >>> with qml.tape.Unwrap(tape1, tape2):
+  ```
   
 <h3>Breaking changes</h3>
 
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 095dc3e59..18464c4e5 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -357,6 +357,17 @@ ar.register_function("torch", "scatter_element_add", _scatter_element_add_torch)
 # -------------------------------- JAX --------------------------------- #
 
 
+def _to_numpy_jax(x):
+    from jax.errors import TracerArrayConversionError
+
+    try:
+        return np.array(getattr(x, "val", x))
+    except TracerArrayConversionError as e:
+        raise ValueError(
+            "Converting a JAX array to a NumPy array not supported when using the JAX JIT."
+        ) from e
+
+
 ar.register_function("jax", "flatten", lambda x: x.flatten())
 ar.register_function(
     "jax",
@@ -364,7 +375,7 @@ ar.register_function(
     lambda x, indices, axis=None: _i("jax").numpy.take(x, indices, axis=axis, mode="wrap"),
 )
 ar.register_function("jax", "coerce", lambda x: x)
-ar.register_function("jax", "to_numpy", lambda x: x)
+ar.register_function("jax", "to_numpy", _to_numpy_jax)
 ar.register_function("jax", "block_diag", lambda x: _i("jax").scipy.linalg.block_diag(*x))
 ar.register_function("jax", "gather", lambda x, indices: x[np.array(indices)])
 ar.register_function(
diff --git a/pennylane/tape/__init__.py b/pennylane/tape/__init__.py
index 017b65c40..a96faa1c9 100644
--- a/pennylane/tape/__init__.py
+++ b/pennylane/tape/__init__.py
@@ -21,3 +21,4 @@ from .cv_param_shift import CVParamShiftTape
 from .qubit_param_shift import QubitParamShiftTape
 from .reversible import ReversibleTape
 from .operation_recorder import OperationRecorder
+from .unwrap import Unwrap, UnwrapTape
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index b0dfa5a8c..4ea8c2f01 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -27,6 +27,8 @@ import pennylane as qml
 from pennylane.queuing import AnnotatedQueue, QueuingContext, QueuingError
 from pennylane.operation import Sample
 
+from .unwrap import UnwrapTape
+
 # CV ops still need to support state preparation operations prior to any
 # other operation for PennyLane-SF tests to pass.
 STATE_PREP_OPS = (
@@ -673,13 +675,8 @@ class QuantumTape(AnnotatedQueue):
         # the current implementation of the adjoint
         # transform requires that the returned inverted object
         # is automatically queued.
-        QuantumTape._lock.acquire()
-        try:
+        with QuantumTape._lock:
             QueuingContext.append(new_tape)
-        except Exception as _:
-            QuantumTape._lock.release()
-            raise
-        QuantumTape._lock.release()
 
         return new_tape
 
@@ -849,6 +846,35 @@ class QuantumTape(AnnotatedQueue):
             op = self._par_info[idx]["op"]
             op.data[self._par_info[idx]["p_idx"]] = p
 
+    def unwrap(self):
+        """A context manager that unwraps a tape with tensor-like parameters
+        to NumPy arrays.
+
+        Args:
+            tape (.QuantumTape): the quantum tape to unwrap
+
+        Returns:
+
+            .QuantumTape: the unwrapped quantum tape
+
+        **Example**
+
+        >>> with tf.GradientTape():
+        ...     with qml.tape.QuantumTape() as tape:
+        ...         qml.RX(tf.Variable(0.1), wires=0)
+        ...         qml.RY(tf.constant(0.2), wires=0)
+        ...         qml.RZ(tf.Variable(0.3), wires=0)
+        ...     with tape.unwrap():
+        ...         print("Trainable params:", tape.trainable_params)
+        ...         print("Unwrapped params:", tape.get_parameters())
+        Trainable params: {0, 2}
+        Unwrapped params: [0.1, 0.3]
+        >>> print("Original parameters:", tape.get_parameters())
+        Original parameters: [<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.1>,
+          <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.3>]
+        """
+        return UnwrapTape(self)
+
     # ========================================================
     # Tape properties
     # ========================================================
diff --git a/pennylane/tape/unwrap.py b/pennylane/tape/unwrap.py
new file mode 100644
index 000000000..8c07ccd23
--- /dev/null
+++ b/pennylane/tape/unwrap.py
@@ -0,0 +1,135 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module contains a context manager for unwrapping tapes
+"""
+import contextlib
+import pennylane as qml
+
+
+class Unwrap:
+    """A context manager that unwraps multiple tapes with tensor-like parameters
+    to NumPy arrays. In addition, this context manager also correctly infers
+    the trainable parameters of the tapes.
+
+    Args:
+        *tapes (.QuantumTape): a sequence of quantum tapes to unwrap
+
+    Returns:
+        Sequence[.QuantumTape]: a sequence of unwrapped quantum tapes
+
+    **Example**
+
+    Consider the following two tapes:
+
+    .. code-block:: python
+
+        x = torch.tensor([0.1, 0.2, 0.3], requires_grad=True, dtype=torch.float64)
+        y = torch.tensor([0.5, 0.6], dtype=torch.float64)
+
+        with qml.tape.QuantumTape() as tape1:
+            qml.RX(x[0], wires=0)
+            qml.RY(y[1], wires=0)
+            qml.RZ(x[2], wires=0)
+
+        with qml.tape.QuantumTape() as tape2:
+            qml.RX(x[1], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.RZ(y[0], wires=0)
+
+    We can use the ``Unwrap`` context manager to simultaneously unwrap the
+    parameters of both tapes:
+
+    >>> with Unwrap(tape1, tape2):
+    ...     print("Tape 1 trainable:", tape1.trainable_params)
+    ...     print("Tape 1 params:", tape1.get_parameters())
+    ...     print("Tape 2 trainable:", tape2.trainable_params)
+    ...     print("Tape 2 params:", tape2.get_parameters())
+    Tape 1 trainable: {0, 2}
+    Tape 1 params: [0.1, 0.3]
+    Tape 2 trainable: {0, 1}
+    Tape 2 params: [0.2, 0.2]
+
+    Outside of the context, the original parameter types remain:
+
+    >>> print("Original parameters:", tape1.get_parameters())
+    Original parameters: [tensor(0.1000, dtype=torch.float64, grad_fn=<SelectBackward>),
+      tensor(0.3000, dtype=torch.float64, grad_fn=<SelectBackward>)]
+    """
+
+    def __init__(self, *tapes):
+        self.tapes = tapes
+        self.stack = None
+
+    def __enter__(self):
+        with contextlib.ExitStack() as stack:
+            for tape in self.tapes:
+                stack.enter_context(UnwrapTape(tape))
+
+            self.stack = stack.pop_all()
+
+        return self.tapes
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self.stack.__exit__(exception_type, exception_value, traceback)
+        self.stack = None
+
+
+class UnwrapTape:
+    """A context manager that unwraps a single tape with tensor-like parameters
+    to NumPy arrays. In addition, this context manager also correctly infers
+    the trainable parameters of the tape.
+
+    Args:
+        tape (.QuantumTape): the quantum tape to unwrap
+
+    Returns:
+
+        .QuantumTape: the unwrapped quantum tape
+
+    **Example**
+
+    >>> with tf.GradientTape():
+    ...     with qml.tape.QuantumTape() as tape:
+    ...         qml.RX(tf.Variable(0.1), wires=0)
+    ...         qml.RY(tf.constant(0.2), wires=0)
+    ...         qml.RZ(tf.Variable(0.3), wires=0)
+    ...     with UnwrapTape(tape) as unwrapped_tape:
+    ...         print("Trainable params:", unwrapped_tape.trainable_params)
+    ...         print("Unwrapped params:", unwrapped_tape.get_parameters())
+    Trainable params: {0, 2}
+    Unwrapped params: [0.1, 0.3]
+    >>> print("Original parameters:", tape.get_parameters())
+    Original parameters: [<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.1>,
+      <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.3>]
+    """
+
+    def __init__(self, tape):
+        self.tape = tape
+        self._original_params = None
+        self._unwrapped_params = None
+
+    def __enter__(self):
+        self._original_params = self.tape.get_parameters(trainable_only=False)
+        self._unwrapped_params = qml.math.unwrap(self._original_params)
+        self.tape.set_parameters(self._unwrapped_params, trainable_only=False)
+
+        # In addition to unwrapping the tape parameters, we also infer the
+        # trainable parameter indices, so that information regarding which
+        # parameters are trainable and which are constant is not lost.
+        self.tape.trainable_params = qml.math.get_trainable_indices(self._original_params)
+        return self.tape
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self.tape.set_parameters(self._original_params, trainable_only=False)


Test patch:
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 7fa2b822a..846009aca 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -610,6 +610,19 @@ def test_numpy(t):
     assert isinstance(res, onp.ndarray)
 
 
+def test_numpy_jax_jit():
+    """Test that the to_numpy() method raises an exception
+    if used inside the JAX JIT"""
+
+    @jax.jit
+    def cost(x):
+        fn.to_numpy(x)
+        return x
+
+    with pytest.raises(ValueError, match="not supported when using the JAX JIT"):
+        cost(jnp.array(0.1))
+
+
 class TestOnesLike:
     """Tests for the ones_like function"""
 
diff --git a/tests/tape/test_unwrap.py b/tests/tape/test_unwrap.py
new file mode 100644
index 000000000..d6a8a947a
--- /dev/null
+++ b/tests/tape/test_unwrap.py
@@ -0,0 +1,252 @@
+# Copyright 2018-2020 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for tape unwrapping"""
+
+import numpy as np
+import pytest
+
+import pennylane as qml
+
+
+def test_unwrap_tensorflow():
+    """Test that unwrapping a tape with TensorFlow parameters
+    works as expected"""
+    tf = pytest.importorskip("tensorflow")
+    p = [tf.Variable(0.1), tf.constant(0.2), np.array(0.5), tf.Variable(0.3)]
+
+    with tf.GradientTape():
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(p[0], wires=0)
+            qml.RY(p[1], wires=0)
+            qml.PhaseShift(p[2], wires=0)
+            qml.RZ(p[3], wires=0)
+
+        with tape.unwrap() as unwrapped_tape:
+            # inside the context manager, all parameters
+            # will be unwrapped to NumPy arrays
+            params = tape.get_parameters(trainable_only=False)
+            assert all(isinstance(i, (float, np.float32)) for i in params)
+            assert np.allclose(params, [0.1, 0.2, 0.5, 0.3])
+            assert tape.trainable_params == {0, 3}
+
+    # outside the context, the original parameters have been restored.
+    assert tape.get_parameters(trainable_only=False) == p
+
+
+def test_unwrap_torch():
+    """Test that unwrapping a tape with Torch parameters
+    works as expected"""
+    torch = pytest.importorskip("torch")
+
+    p = [
+        torch.tensor(0.1, requires_grad=True),
+        torch.tensor(0.2),
+        np.array(0.5),
+        torch.tensor(0.3, requires_grad=True),
+    ]
+
+    with qml.tape.QuantumTape() as tape:
+        qml.RX(p[0], wires=0)
+        qml.RY(p[1], wires=0)
+        qml.PhaseShift(p[2], wires=0)
+        qml.RZ(p[3], wires=0)
+
+    with tape.unwrap() as unwrapped_tape:
+        # inside the context manager, all parameters
+        # will be unwrapped to NumPy arrays
+        params = tape.get_parameters(trainable_only=False)
+        assert all(isinstance(i, float) for i in params)
+        assert np.allclose(params, [0.1, 0.2, 0.5, 0.3])
+        assert tape.trainable_params == {0, 3}
+
+    # outside the context, the original parameters have been restored.
+    assert tape.get_parameters(trainable_only=False) == p
+
+
+def test_unwrap_autograd():
+    """Test that unwrapping a tape with Autograd parameters
+    works as expected"""
+    from pennylane import numpy as anp
+
+    p = [
+        anp.tensor(0.1, requires_grad=True),
+        anp.tensor(0.2, requires_grad=False),
+        0.5,
+        anp.tensor(0.3, requires_grad=True),
+    ]
+
+    with qml.tape.QuantumTape() as tape:
+        qml.RX(p[0], wires=0)
+        qml.RY(p[1], wires=0)
+        qml.PhaseShift(p[2], wires=0)
+        qml.RZ(p[3], wires=0)
+
+    with tape.unwrap() as unwrapped_tape:
+        # inside the context manager, all parameters
+        # will be unwrapped to NumPy arrays
+        params = tape.get_parameters(trainable_only=False)
+        assert all(isinstance(i, float) for i in params)
+        assert np.allclose(params, [0.1, 0.2, 0.5, 0.3])
+        assert tape.trainable_params == {0, 2, 3}
+
+    # outside the context, the original parameters have been restored.
+    assert tape.get_parameters(trainable_only=False) == p
+
+
+def test_unwrap_autograd_backward():
+    """Test that unwrapping a tape with Autograd parameters
+    works as expected during a backwards pass"""
+    from pennylane import numpy as anp
+    from autograd.numpy.numpy_boxes import ArrayBox
+
+    p = [
+        anp.tensor([0.1, 0.5, 0.3], requires_grad=True),
+        anp.tensor(0.2, requires_grad=False),
+    ]
+
+    def cost(*p):
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(p[0][0], wires=0)
+            qml.RY(p[1], wires=0)
+            qml.PhaseShift(p[0][1], wires=0)
+            qml.RZ(p[0][2], wires=0)
+
+        with tape.unwrap() as unwrapped_tape:
+            # inside the context manager, all parameters
+            # will be unwrapped to NumPy arrays
+            params = tape.get_parameters(trainable_only=False)
+            assert all(isinstance(i, float) for i in params)
+            assert np.allclose(params, [0.1, 0.2, 0.5, 0.3])
+            assert tape.trainable_params == {0, 2, 3}
+
+        # outside the context, the original parameters have been restored.
+        params = tape.get_parameters(trainable_only=False)
+        assert any(isinstance(i, ArrayBox) for i in params)
+
+        return p[0][0] * p[1] ** 2 * anp.sin(p[0][1]) * anp.exp(-0.5 * p[0][2])
+
+    qml.grad(cost)(*p)
+    qml.jacobian(qml.grad(cost))(*p)
+
+
+def test_unwrap_jax():
+    """Test that unwrapping a tape with JAX parameters
+    works as expected"""
+    jax = pytest.importorskip("jax")
+    from jax import numpy as jnp
+
+    p = [
+        jnp.array(0.1),
+        jnp.array(0.2),
+        0.5,
+        jnp.array(0.3),
+    ]
+
+    with qml.tape.QuantumTape() as tape:
+        qml.RX(p[0], wires=0)
+        qml.RY(p[1], wires=0)
+        qml.PhaseShift(p[2], wires=0)
+        qml.RZ(p[3], wires=0)
+
+    with tape.unwrap() as unwrapped_tape:
+        # inside the context manager, all parameters
+        # will be unwrapped to NumPy arrays
+        params = tape.get_parameters(trainable_only=False)
+        assert all(isinstance(i, float) for i in params)
+        assert np.allclose(params, [0.1, 0.2, 0.5, 0.3])
+
+        # During the forward pass, JAX has no concept of trainable
+        # arrays.
+        assert tape.trainable_params == set()
+
+    # outside the context, the original parameters have been restored.
+    assert tape.get_parameters(trainable_only=False) == p
+
+
+def test_unwrap_jax_backward():
+    """Test that unwrapping a tape with JAX parameters
+    works as expected during a backwards pass"""
+    jax = pytest.importorskip("jax")
+    from jax import numpy as jnp
+    from jaxlib.xla_extension import DeviceArray
+    from jax.interpreters.ad import JVPTracer
+
+    p = [
+        jnp.array([0.1, 0.5, 0.3]),
+        jnp.array(0.2),
+    ]
+
+    def cost(*p):
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(p[0][0], wires=0)
+            qml.RY(p[1], wires=0)
+            qml.PhaseShift(p[0][1], wires=0)
+            qml.RZ(p[0][2], wires=0)
+
+        with tape.unwrap() as unwrapped_tape:
+            # inside the context manager, all parameters
+            # will be unwrapped to NumPy arrays
+            params = tape.get_parameters(trainable_only=False)
+            assert all(isinstance(i, float) for i in params)
+            assert np.allclose(params, [0.1, 0.2, 0.5, 0.3])
+            assert tape.trainable_params == {0, 2, 3}
+
+        # outside the context, the original parameters have been restored.
+        params = tape.get_parameters(trainable_only=False)
+        assert all(isinstance(i, (DeviceArray, JVPTracer)) for i in params)
+
+        return p[0][0] * p[1] ** 2 * jnp.sin(p[0][1]) * jnp.exp(-0.5 * p[0][2])
+
+    jax.grad(cost, argnums=0)(*p)
+    jax.jacobian(jax.grad(cost, argnums=0), argnums=0)(*p)
+
+
+def test_multiple_unwrap():
+    """Test that unwrapping multiple tapes at once works correctly"""
+    torch = pytest.importorskip("torch")
+
+    p = [
+        torch.tensor(0.1, requires_grad=True),
+        torch.tensor(0.2),
+        np.array(0.5),
+        torch.tensor(0.3, requires_grad=True),
+    ]
+
+    with qml.tape.QuantumTape() as tape1:
+        qml.RX(p[0], wires=0)
+        qml.RY(p[1], wires=0)
+        qml.PhaseShift(p[2], wires=0)
+        qml.RZ(p[3], wires=0)
+
+    with qml.tape.QuantumTape() as tape2:
+        qml.RX(p[1], wires=0)
+        qml.RY(p[3], wires=0)
+        qml.PhaseShift(p[0], wires=0)
+        qml.RZ(p[2], wires=0)
+
+    with qml.tape.Unwrap(tape1, tape2):
+        # inside the context manager, all parameters
+        # will be unwrapped to NumPy arrays
+        params = tape1.get_parameters(trainable_only=False)
+        assert all(isinstance(i, float) for i in params)
+        assert tape1.trainable_params == {0, 3}
+
+        params = tape2.get_parameters(trainable_only=False)
+        assert all(isinstance(i, float) for i in params)
+        assert tape2.trainable_params == {1, 2}
+
+    # outside the context, the original parameters have been restored.
+    assert tape1.get_parameters(trainable_only=False) == p
+    assert tape2.get_parameters(trainable_only=False) == [p[1], p[3], p[0], p[2]]