fallback to scipy when linalg.matrix_power isn't available

doc/releases/changelog-dev.md

@@ -106,6 +106,10 @@
   expected, instead of numpy arrays.
   [(#4811)](https://github.com/PennyLaneAI/pennylane/pull/4811)
 
+* `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_pow_op.py

@@ -744,6 +744,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."""
 

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)]