Repository: PennyLaneAI/pennylane
Pull request #5509 (2024-04-12 20:01:20)

Add `qml.capture` module
**Context:**
A new program capturing mechanism within PennyLane is planned.

**Description of the Change:**
This PR adds the `qml.capture` module which will contain code for said mechanism.

**Benefits:**
Module

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**

[sc-59321]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5509?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.66%. Comparing base [(`b2b0c36`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b2b0c36608cce4772cda7200d776d67eac9dfff6?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`f70f29b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5509?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 9 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5509      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         404      406       +2     
  Lines       37833    37577     -256     
==========================================
- Hits        37709    37452     -257     
- Misses        124      125       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5509?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/code/qml_capture.rst b/doc/code/qml_capture.rst
new file mode 100644
index 000000000..0d957f80d
--- /dev/null
+++ b/doc/code/qml_capture.rst
@@ -0,0 +1,7 @@
+qml.capture
+===========
+
+.. currentmodule:: pennylane.capture
+
+.. automodule:: pennylane.capture
+
diff --git a/doc/code/qml_compiler.rst b/doc/code/qml_compiler.rst
index e5bb5ee63..1b56df578 100644
--- a/doc/code/qml_compiler.rst
+++ b/doc/code/qml_compiler.rst
@@ -1,5 +1,5 @@
 qml.compiler
-===============
+============
 
 .. currentmodule:: pennylane.compiler
 
diff --git a/doc/index.rst b/doc/index.rst
index b70027195..e03e6068a 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -213,6 +213,7 @@ PennyLane is **free** and **open source**, released under the Apache License, Ve
    :caption: Internals
    :hidden:
 
+   code/qml_capture
    code/qml_devices
    code/qml_measurements
    code/qml_operation
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index faebb6e6b..fea25758b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* Added a qml.capture module that will contain PennyLane's own capturing mechanism for hybrid
+  quantum-classical programs.
+  [(#5509)](https://github.com/PennyLaneAI/pennylane/pull/5509)
+
 * The `FABLE` template is added for efficient block encoding of matrices. Users can now call FABLE to efficiently construct circuits according to a user-set approximation level. 
 [(#5107)](https://github.com/PennyLaneAI/pennylane/pull/5107)
 
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 081d214fb..85bb1b7b0 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -26,6 +26,7 @@ from pennylane.boolean_fn import BooleanFn
 import pennylane.numpy
 from pennylane.queuing import QueuingManager, apply
 
+import pennylane.capture
 import pennylane.kernels
 import pennylane.math
 import pennylane.operation
diff --git a/pennylane/capture/__init__.py b/pennylane/capture/__init__.py
new file mode 100644
index 000000000..24feeaebf
--- /dev/null
+++ b/pennylane/capture/__init__.py
@@ -0,0 +1,44 @@
+# Copyright 2024 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.
+r"""
+.. currentmodule:: pennylane
+
+This module implements PennyLane's capturing mechanism for hybrid
+quantum-classical programs.
+
+.. warning::
+
+    This module is experimental and will change significantly in the future.
+
+
+To activate and deactivate the new PennyLane program capturing mechanism, use
+the switches ``qml.capture.enable_plxpr`` and ``qml.capture.disable_plxpr``.
+Whether or not the capturing mechanism is currently being used can be
+queried with ``qml.capture.plxpr_enabled``.
+By default, the mechanism is disabled:
+
+.. code-block:: pycon
+
+    >>> import pennylane as qml
+    >>> qml.capture.plxpr_enabled()
+    False
+    >>> qml.capture.enable_plxpr()
+    >>> qml.capture.plxpr_enabled()
+    True
+    >>> qml.capture.disable_plxpr()
+    >>> qml.capture.plxpr_enabled()
+    False
+
+"""
+from .switches import enable_plxpr, disable_plxpr, plxpr_enabled
diff --git a/pennylane/capture/switches.py b/pennylane/capture/switches.py
new file mode 100644
index 000000000..0b32e6d37
--- /dev/null
+++ b/pennylane/capture/switches.py
@@ -0,0 +1,64 @@
+# Copyright 2024 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.
+"""
+Contains the switches to (de)activate the capturing mechanism, and a
+status reporting function on whether it is enabled or not.
+"""
+from typing import Callable
+
+has_jax = True
+try:
+    import jax  # pylint: disable=unused-import
+except ImportError:
+    has_jax = False
+
+
+def _make_switches() -> [Callable[[], None], Callable[[], None], Callable[[], bool]]:
+    r"""Create three functions, corresponding to an activation switch, a deactivation switch
+    and a status query, in that order.
+
+    .. note::
+
+        While the internal variable is named in some context, this function
+        can be used to make switches for any context.
+    """
+
+    _PLXPR_ENABLED = False
+    # since this changes what happens with tracing, we need to turn the behavior
+    # off by default to preserve our ability to jit pennylane circuits.
+
+    def enable() -> None:
+        """Enable the capturing mechanism of hybrid quantum-classical programs
+        in a PennyLane Program Representation (plxpr)."""
+        if not has_jax:
+            raise ImportError("plxpr requires JAX to be installed.")
+        nonlocal _PLXPR_ENABLED
+        _PLXPR_ENABLED = True
+
+    def disable() -> None:
+        """Disable the capturing mechanism of hybrid quantum-classical programs
+        in a PennyLane Program Representation (plxpr)."""
+        nonlocal _PLXPR_ENABLED
+        _PLXPR_ENABLED = False
+
+    def status() -> bool:
+        """Return whether the capturing mechanism of hybrid quantum-classical programs
+        in a PennyLane Program Representation (plxpr) is enabled."""
+        nonlocal _PLXPR_ENABLED
+        return _PLXPR_ENABLED
+
+    return enable, disable, status
+
+
+enable_plxpr, disable_plxpr, plxpr_enabled = _make_switches()


Test patch:
diff --git a/tests/capture/test_switches.py b/tests/capture/test_switches.py
new file mode 100644
index 000000000..2d0310c49
--- /dev/null
+++ b/tests/capture/test_switches.py
@@ -0,0 +1,40 @@
+# Copyright 2024 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 the backbone of the :mod:`pennylane.capture` module.
+"""
+import pytest
+import pennylane as qml
+
+
+@pytest.mark.jax
+def test_switches_with_jax():
+    """Test switches and status reporting function."""
+
+    assert qml.capture.plxpr_enabled() is False
+    assert qml.capture.enable_plxpr() is None
+    assert qml.capture.plxpr_enabled() is True
+    assert qml.capture.disable_plxpr() is None
+    assert qml.capture.plxpr_enabled() is False
+
+
+def test_switches_without_jax():
+    """Test switches and status reporting function."""
+
+    assert qml.capture.plxpr_enabled() is False
+    with pytest.raises(ImportError, match="plxpr requires JAX to be installed."):
+        qml.capture.enable_plxpr()
+    assert qml.capture.plxpr_enabled() is False
+    assert qml.capture.disable_plxpr() is None
+    assert qml.capture.plxpr_enabled() is False
diff --git a/tests/test_tensor_measurements.py b/tests/test_tensor_measurements.py
index bc6386251..9fd6a7b03 100644
--- a/tests/test_tensor_measurements.py
+++ b/tests/test_tensor_measurements.py
@@ -53,7 +53,7 @@ class TestTensorExpval:
     def test_tensor_product(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product ZxZ gives the same result as simply
         using an Hermitian matrix"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=1851)
 
         @qml.qnode(dev)
         def circuit1(a, b, c):
@@ -73,7 +73,7 @@ class TestTensorExpval:
     def test_combine_tensor_with_non_tensor(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product along with a non-tensor product
         continues to function correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=181)
 
         @qml.qnode(dev)
         def circuit1(a, b, c):
@@ -97,7 +97,7 @@ class TestTensorExpval:
 
     def test_paulix_tensor_pauliy(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=164)
 
         @qml.qnode(dev)
         def circuit(a, b, c):
@@ -112,7 +112,7 @@ class TestTensorExpval:
     @pytest.mark.autograd
     def test_paulix_tensor_pauliy_gradient(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=144)
 
         @qml.qnode(dev)
         def circuit(a, b, c):
@@ -127,7 +127,7 @@ class TestTensorExpval:
 
     def test_pauliz_tensor_identity(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving PauliZ and Identity works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=134)
 
         @qml.qnode(dev)
         def circuit(a, b, c):
@@ -141,7 +141,7 @@ class TestTensorExpval:
 
     def test_pauliz_tensor_hadamard(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving PauliZ and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=324)
 
         @qml.qnode(dev)
         def circuit(a, b, c):
@@ -155,7 +155,7 @@ class TestTensorExpval:
 
     def test_hermitian(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving an Hermitian matrix works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=125)
 
         A = np.array(
             [
@@ -183,7 +183,7 @@ class TestTensorExpval:
 
     def test_hermitian_tensor_hermitian(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving two Hermitian matrices works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=224)
 
         A1 = np.array([[1, 2], [2, 4]])
 
@@ -223,7 +223,7 @@ class TestTensorExpval:
 
     def test_hermitian_tensor_identity_expectation(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving an Hermitian matrix and the identity works correctly"""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit", wires=2, shots=shots, seed=144)
 
         A = np.array(
             [[1.02789352, 1.61296440 - 0.3498192j], [1.61296440 + 0.3498192j, 1.23920938 + 0j]]
@@ -261,7 +261,7 @@ class TestTensorVar:
 
     def test_paulix_tensor_pauliy(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=924)
 
         @qml.qnode(dev)
         def circuit(a, b, c):
@@ -282,7 +282,7 @@ class TestTensorVar:
 
     def test_pauliz_tensor_hadamard(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving PauliZ and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=167)
 
         @qml.qnode(dev)
         def circuit(a, b, c):
@@ -301,7 +301,7 @@ class TestTensorVar:
 
     def test_tensor_hermitian(self, shots, theta, phi, varphi, tolerance):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots, seed=824)
 
         A = np.array(
             [
@@ -363,7 +363,7 @@ class TestTensorSample:
     # pylint: disable=unused-argument
     def test_paulix_tensor_pauliz(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliX and PauliZ works correctly"""
-        dev = qml.device("default.qubit", wires=2, shots=int(1e6))
+        dev = qml.device("default.qubit", wires=2, shots=int(1e6), seed=524)
 
         @qml.qnode(dev)
         def circuit():
@@ -377,7 +377,7 @@ class TestTensorSample:
 
     def test_paulix_tensor_pauliy(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit", wires=3, shots=int(1e6), seed=173)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(a, b, c):
@@ -391,7 +391,7 @@ class TestTensorSample:
 
     def test_pauliz_tensor_hadamard(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliZ and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit", wires=3, shots=int(1e6), seed=814)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(a, b, c):
@@ -405,7 +405,7 @@ class TestTensorSample:
 
     def test_tensor_hermitian(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit", wires=3, shots=int(1e6), seed=124)
 
         A = np.array(
             [