Exp matrix computation bugfix (#3057)

* Incrementing the version number to `v0.27.0-dev` (#3044)

* Update

* Bump

* re-add dev changelog

* Update doc/development/release_notes.md

* docs section

* flaky test_single_argument_step

Co-authored-by: Antal Szava <antalszava@gmail.com>

* fix exp bug

* rebase onto release candidate

* Update doc/releases/changelog-0.26.0.md

* Update tests/optimize/test_optimize_shot_adapative.py

* add additional test parametrization

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
Co-authored-by: Antal Szava <antalszava@gmail.com>

.pre-commit-config.yaml

@@ -18,4 +18,3 @@ repos:
     hooks:
     - id: black
       exclude: ^doc/
-      language_version: python3.8

doc/releases/changelog-0.26.0.md

@@ -491,6 +491,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug with `qml.ops.Exp` operators when the coefficient is autograd but the diagonalizing gates
+  don't act on all wires.
+  [(#3057)](https://github.com/PennyLaneAI/pennylane/pull/3057)
+
 * Fixed a bug where the tape transform `single_qubit_fusion` computed wrong rotation angles
   for specific combinations of rotations.
   [(#3024)](https://github.com/PennyLaneAI/pennylane/pull/3024)

pennylane/ops/op_math/exp.py

@@ -168,14 +168,17 @@ def _queue_category(self):
     def matrix(self, wire_order=None):
 
         coeff_interface = math.get_interface(self.coeff)
-        if coeff_interface == "autograd":
+        if coeff_interface == "autograd" and math.requires_grad(self.coeff):
             # math.expm is not differentiable with autograd
             # So we try to do a differentiable construction if possible
             #
             # This won't catch situations when the base matrix is autograd,
             # but at least this provides as much trainablility as possible
             try:
-                diagonalizing_mat = qml.matrix(self.diagonalizing_gates)()
+                if len(self.diagonalizing_gates()) == 0:
+                    eigvals_mat = math.diag(self.eigvals())
+                    return expand_matrix(eigvals_mat, wires=self.wires, wire_order=wire_order)
+                diagonalizing_mat = qml.matrix(self.diagonalizing_gates, wire_order=self.wires)()
                 eigvals_mat = math.diag(self.eigvals())
                 mat = diagonalizing_mat.conj().T @ eigvals_mat @ diagonalizing_mat
                 return expand_matrix(mat, wires=self.wires, wire_order=wire_order)

tests/ops/op_math/test_exp.py

@@ -139,14 +139,38 @@ def test_prod_base_isingyy(self):
         assert qml.math.allclose(op.matrix(), isingxx.matrix())
 
     @pytest.mark.autograd
-    def test_matrix_autograd_rx(self):
+    @pytest.mark.parametrize("requires_grad", (True, False))
+    def test_matrix_autograd_rx(self, requires_grad):
         """Test the matrix comparing to the rx gate."""
-        phi = np.array(1.234)
+        phi = np.array(1.234, requires_grad=requires_grad)
         exp_rx = Exp(qml.PauliX(0), -0.5j * phi)
         rx = qml.RX(phi, 0)
 
         assert qml.math.allclose(exp_rx.matrix(), rx.matrix())
 
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("requires_grad", (True, False))
+    def test_matrix_autograd_rz(self, requires_grad):
+        """Test the matrix comparing to the rz gate. This is a gate with an
+        autograd coefficient but empty diagonalizing gates."""
+        phi = np.array(1.234, requires_grad=requires_grad)
+        exp_rz = Exp(qml.PauliZ(0), -0.5j * phi)
+        rz = qml.RZ(phi, 0)
+
+        assert qml.math.allclose(exp_rz.matrix(), rz.matrix())
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("requires_grad", (True, False))
+    def test_tensor_with_pauliz_autograd(self, requires_grad):
+        """Test the matrix for the case when the coefficient is autograd and
+        the diagonalizing gates don't act on every wire for the matrix."""
+        phi = qml.numpy.array(-0.345, requires_grad=requires_grad)
+        base = qml.PauliZ(0) @ qml.PauliY(1)
+        autograd_op = Exp(base, phi)
+        mat = qml.math.expm(phi * qml.matrix(base))
+
+        assert qml.math.allclose(autograd_op.matrix(), mat)
+
     @pytest.mark.autograd
     def test_base_no_diagonalizing_gates_autograd_coeff(self):
         """Test the matrix when the base matrix doesn't define the diagonalizing gates."""