fallback to scipy when linalg.matrix_power isn't available

doc/releases/changelog-dev.md

@@ -128,6 +128,10 @@
   with QNodes.
   [(#4831)](https://github.com/PennyLaneAI/pennylane/pull/4831)
 
+* `qml.pow(op)` and `qml.QubitUnitary.pow()` now also work with Tensorflow data raised to an
+  integer power.
+  [(#4827)](https://github.com/PennyLaneAI/pennylane/pull/4827)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):

pennylane/ops/op_math/pow.py

@@ -259,7 +259,7 @@ def label(self, decimals=None, base_label=None, cache=None):
 
     @staticmethod
     def _matrix(scalar, mat):
-        if isinstance(scalar, int):
+        if isinstance(scalar, int) and qml.math.get_deep_interface(mat) != "tensorflow":
             return qmlmath.linalg.matrix_power(mat, scalar)
         return fractional_matrix_power(mat, scalar)
 

pennylane/ops/qubit/matrix_ops.py

@@ -20,6 +20,7 @@
 from itertools import product
 
 import numpy as np
+from scipy.linalg import fractional_matrix_power
 from pennylane.math import norm, cast, eye, zeros, transpose, conj, sqrt, sqrt_matrix
 from pennylane import numpy as pnp
 
@@ -232,9 +233,14 @@ def adjoint(self):
         return QubitUnitary(qml.math.moveaxis(qml.math.conj(U), -2, -1), wires=self.wires)
 
     def pow(self, z):
-        if isinstance(z, int):
-            return [QubitUnitary(qml.math.linalg.matrix_power(self.matrix(), z), wires=self.wires)]
-        return super().pow(z)
+        mat = self.matrix()
+        if isinstance(z, int) and qml.math.get_deep_interface(mat) != "tensorflow":
+            pow_mat = qml.math.linalg.matrix_power(mat, z)
+        elif self.batch_size is not None or qml.math.shape(z) != ():
+            return super().pow(z)
+        else:
+            pow_mat = qml.math.convert_like(fractional_matrix_power(mat, z), mat)
+        return [QubitUnitary(pow_mat, wires=self.wires)]
 
     def _controlled(self, wire):
         return qml.ControlledQubitUnitary(*self.parameters, control_wires=wire, wires=self.wires)

tests/ops/op_math/test_controlled_ops.py

@@ -19,6 +19,7 @@
 
 import numpy as np
 import pytest
+from scipy.linalg import fractional_matrix_power
 from scipy.sparse import csr_matrix
 from scipy.stats import unitary_group
 
@@ -393,7 +394,7 @@ def test_pow_broadcasted(self, n):
         assert qml.math.allclose(pow_ops[0].data[0], op_mat_to_pow)
 
     def test_noninteger_pow(self):
-        """Test that a ControlledQubitUnitary raised to a non-integer power raises an error."""
+        """Test that a ControlledQubitUnitary raised to a non-integer power evalutes."""
         U1 = np.array(
             [
                 [0.73708696 + 0.61324932j, 0.27034258 + 0.08685028j],
@@ -403,8 +404,11 @@ def test_noninteger_pow(self):
 
         op = qml.ControlledQubitUnitary(U1, control_wires=("b", "c"), wires="a")
 
-        with pytest.raises(qml.operation.PowUndefinedError):
-            op.pow(0.12)
+        z = 0.12
+        [pow_op] = op.pow(z)
+        expected = np.eye(8, dtype=complex)
+        expected[-2:, -2:] = fractional_matrix_power(U1, z)
+        assert qml.math.allequal(pow_op.matrix(), expected)
 
     def test_noninteger_pow_broadcasted(self):
         """Test that a ControlledQubitUnitary raised to a non-integer power raises an error."""

tests/ops/op_math/test_pow_op.py

@@ -706,9 +706,9 @@ def check_matrix(self, param, z):
         base = qml.IsingZZ(param, wires=(0, 1))
         op = Pow(base, z)
 
-        mat = qml.matrix(op)
-        shortcut = base.pow(z)[0]
-        shortcut_mat = qml.matrix(shortcut)
+        mat = op.matrix()
+        [shortcut] = base.pow(z)
+        shortcut_mat = shortcut.matrix()
 
         return qml.math.allclose(mat, shortcut_mat)
 
@@ -751,6 +751,15 @@ def test_matrix_against_shortcut_tf(self, z):
         param = tf.Variable(2.34)
         assert self.check_matrix(param, z)
 
+    @pytest.mark.tf
+    def test_matrix_tf_int_z(self):
+        """Test that matrix works with integer power."""
+        import tensorflow as tf
+
+        theta = tf.Variable(1.0)
+        mat = qml.pow(qml.RX(theta, wires=0), z=3).matrix()
+        assert qml.math.allclose(mat, qml.RX.compute_matrix(3))
+
     def test_matrix_wire_order(self):
         """Test that the wire_order keyword rearranges ording."""
 
@@ -796,7 +805,7 @@ def test_sparse_matrix_exists_int_exponent(self):
         sparse_mat_array = sparse_mat.toarray()
 
         assert qml.math.allclose(sparse_mat_array, H_cubed.toarray())
-        assert qml.math.allclose(sparse_mat_array, qml.matrix(op))
+        assert qml.math.allclose(sparse_mat_array, op.matrix())
 
     def test_sparse_matrix_float_exponent(self):
         """Test that even a sparse-matrix defining op raised to a float power

tests/ops/qubit/test_matrix_ops.py

@@ -25,7 +25,7 @@
 from pennylane import numpy as pnp
 from pennylane.operation import DecompositionUndefinedError
 from pennylane.wires import Wires
-from pennylane.ops.qubit.matrix_ops import _walsh_hadamard_transform
+from pennylane.ops.qubit.matrix_ops import _walsh_hadamard_transform, fractional_matrix_power
 
 
 class TestQubitUnitary:
@@ -38,9 +38,10 @@ def test_qubit_unitary_noninteger_pow(self):
         )
 
         op = qml.QubitUnitary(U, wires="a")
+        [pow_op] = op.pow(0.123)
+        expected = fractional_matrix_power(U, 0.123)
 
-        with pytest.raises(qml.operation.PowUndefinedError):
-            op.pow(0.123)
+        assert qml.math.allclose(pow_op.matrix(), expected)
 
     def test_qubit_unitary_noninteger_pow_broadcasted(self):
         """Test broadcasted QubitUnitary raised to a non-integer power raises an error."""
@@ -188,6 +189,16 @@ def test_qubit_unitary_tf(self, U, num_wires):
         with pytest.raises(ValueError, match="must be of shape"):
             qml.QubitUnitary(U, wires=range(num_wires + 1)).matrix()
 
+    @pytest.mark.tf
+    def test_qubit_unitary_int_pow_tf(self):
+        """Test that QubitUnitary.pow works with tf and int z values."""
+        import tensorflow as tf
+
+        mat = tf.Variable([[1, 0], [0, tf.exp(1j)]])
+        expected = tf.Variable([[1, 0], [0, tf.exp(3j)]])
+        [op] = qml.QubitUnitary(mat, wires=[0]).pow(3)
+        assert qml.math.allclose(op.matrix(), expected)
+
     @pytest.mark.jax
     @pytest.mark.parametrize(
         "U,num_wires", [(H, 1), (np.kron(H, H), 2), (np.tensordot([1j, -1, 1], H, axes=0), 1)]