Repository: PennyLaneAI/pennylane
Pull request #3327 (2022-11-11 15:02:40)

CMP #2: Add `MutualInfo` MP class
1. Create `_MutualInfo` class (name will be changed to `MutualInfo` after the whole refactor).
2. Return `_MutualInfo` instance in `qml.mutual_info`.
3. Copy code from `QubitDevice` into `_MutualInfo.process_state`.
4. Added test.


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8f77437f4..2adb42948 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -9,6 +9,9 @@
     method to process samples/quantum state.
     [#3286](https://github.com/PennyLaneAI/pennylane/pull/3286)
 
+  * Add `_MutualInfo` class.
+    [#3327](https://github.com/PennyLaneAI/pennylane/pull/3327)
+
 * Functionality for fetching symbols and geometry of a compound from the PubChem Database using `qchem.mol_data`.
   [(#3289)](https://github.com/PennyLaneAI/pennylane/pull/3289)
  
diff --git a/pennylane/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py
index 2d494a99a..8474c5804 100644
--- a/pennylane/measurements/mutual_info.py
+++ b/pennylane/measurements/mutual_info.py
@@ -15,9 +15,12 @@
 """
 This module contains the qml.mutual_info measurement.
 """
+from typing import Sequence
+
 import pennylane as qml
+from pennylane.wires import Wires
 
-from .measurements import MeasurementProcess, MutualInfo
+from .measurements import MutualInfo, StateMeasurement
 
 
 def mutual_info(wires0, wires1, log_base=None):
@@ -76,4 +79,17 @@ def mutual_info(wires0, wires1, log_base=None):
 
     wires0 = qml.wires.Wires(wires0)
     wires1 = qml.wires.Wires(wires1)
-    return MeasurementProcess(MutualInfo, wires=[wires0, wires1], log_base=log_base)
+    return _MutualInfo(MutualInfo, wires=[wires0, wires1], log_base=log_base)
+
+
+class _MutualInfo(StateMeasurement):
+    """Measurement process that returns the mutual information."""
+
+    def process_state(self, state: Sequence[complex], wires: Wires):
+        return qml.math.mutual_info(
+            state,
+            indices0=list(self._wires[0]),
+            indices1=list(self._wires[1]),
+            c_dtype=state.dtype,
+            base=self.log_base,
+        )


Test patch:
diff --git a/tests/measurements/test_mutual_info.py b/tests/measurements/test_mutual_info.py
new file mode 100644
index 000000000..142c46fc9
--- /dev/null
+++ b/tests/measurements/test_mutual_info.py
@@ -0,0 +1,65 @@
+# 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 the mutual_info module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.interfaces import INTERFACE_MAP
+from pennylane.measurements.mutual_info import _MutualInfo
+
+
+class TestMutualInfo:
+    """Tests for the mutual_info function"""
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
+    @pytest.mark.parametrize(
+        "state, expected",
+        [
+            ([1.0, 0.0, 0.0, 0.0], 0),
+            ([qml.math.sqrt(2) / 2, 0.0, qml.math.sqrt(2) / 2, 0.0], 0),
+            ([qml.math.sqrt(2) / 2, 0.0, 0.0, qml.math.sqrt(2) / 2], 2 * qml.math.log(2)),
+            (qml.math.ones(4) * 0.5, 0.0),
+        ],
+    )
+    def test_mutual_info_output(self, interface, state, expected):
+        """Test the output of qml.mutual_info"""
+        dev = qml.device("default.qubit", wires=4)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            qml.QubitStateVector(state, wires=[0, 1])
+            return qml.mutual_info(wires0=[0, 2], wires1=[1, 3])
+
+        res = circuit()
+        new_res = qml.mutual_info(wires0=[0, 2], wires1=[1, 3]).process_state(
+            state=circuit.device.state, wires=circuit.device.wires
+        )
+        assert np.allclose(res, expected, atol=1e-6)
+        assert np.allclose(new_res, expected, atol=1e-6)
+        assert INTERFACE_MAP.get(qml.math.get_interface(new_res)) == interface
+        assert res.dtype == new_res.dtype
+
+    def test_queue(self):
+        """Test that the right measurement class is queued."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        circuit()
+
+        assert isinstance(circuit.tape[0], _MutualInfo)