Improve computation time of compute_matrix for parametric ops

doc/releases/changelog-dev.md

@@ -40,12 +40,18 @@
   -1.1258709813834058
   ```
 
-* New PennyLane-inspired `sketch` and `sketch_dark` styles are now available for drawing circuit diagram graphics. 
+* New PennyLane-inspired `sketch` and `sketch_dark` styles are now available for
+  drawing circuit diagram graphics. 
   [(#2709)](https://github.com/PennyLaneAI/pennylane/pull/2709)
 
 <h3>Improvements</h3>
 
-* Adds a new function to compare operators. `qml.equal` can be used to compare equality of parametric operators taking into account their interfaces and trainability.
+* A small performance and readability upgrade to the `compute_matrix` method
+  of broadcastable parametric operations.
+  [(#2726)](https://github.com/PennyLaneAI/pennylane/pull/2726)
+
+* Adds a new function to compare operators. `qml.equal` can be used to compare equality of
+  parametric operators taking into account their interfaces and trainability.
   [(#2651)](https://github.com/PennyLaneAI/pennylane/pull/2651)
 
 <h3>Breaking changes</h3>
@@ -58,4 +64,4 @@
 
 This release contains contributions from (in alphabetical order):
 
-Ankit Khandelwal, Ixchel Meza Chavez, Moritz Willmann
+Ankit Khandelwal, Ixchel Meza Chavez, David Wierichs, Moritz Willmann

pennylane/ops/qubit/parametric_ops.py

@@ -276,10 +276,12 @@ def compute_matrix(theta):  # pylint: disable=arguments-differ
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)
 
-        p = qml.math.exp(-0.5j * theta)
-        z = qml.math.zeros_like(p)
+        arg = 0.5j * theta
+        if qml.math.ndim(arg) == 0:
+            return qml.math.diag(qml.math.exp(arg * np.array([-1, 1])))
 
-        return qml.math.stack([stack_last([p, z]), stack_last([z, qml.math.conj(p)])], axis=-2)
+        diags = qml.math.exp(qml.math.outer(arg, [-1, 1]))
+        return diags[:, :, np.newaxis] * np.eye(2)
 
     @staticmethod
     def compute_eigvals(theta):  # pylint: disable=arguments-differ
@@ -312,9 +314,10 @@ def compute_eigvals(theta):  # pylint: disable=arguments-differ
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)
 
-        p = qml.math.exp(-0.5j * theta)
+        if qml.math.ndim(theta) == 0:
+            return qml.math.exp(0.5j * theta * np.array([-1, 1]))
 
-        return stack_last([p, qml.math.conj(p)])
+        return qml.math.exp(qml.math.outer(0.5j * theta, [-1, 1]))
 
     def adjoint(self):
         return RZ(-self.data[0], wires=self.wires)
@@ -399,10 +402,12 @@ def compute_matrix(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        p = qml.math.exp(1j * phi)
-        z = qml.math.zeros_like(p)
+        arg = 1j * phi
+        if qml.math.ndim(arg) == 0:
+            return qml.math.diag(qml.math.exp(arg * np.array([0, 1])))
 
-        return qml.math.stack([stack_last([qml.math.ones_like(p), z]), stack_last([z, p])], axis=-2)
+        diags = qml.math.exp(qml.math.outer(arg, [0, 1]))
+        return diags[:, :, np.newaxis] * np.eye(2)
 
     @staticmethod
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
@@ -434,9 +439,10 @@ def compute_eigvals(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        p = qml.math.exp(1j * phi)
+        if qml.math.ndim(phi) == 0:
+            return qml.math.exp(1j * phi * np.array([0, 1]))
 
-        return stack_last([qml.math.ones_like(p), p])
+        return qml.math.exp(qml.math.outer(1j * phi, [0, 1]))
 
     @staticmethod
     def compute_decomposition(phi, wires):
@@ -550,21 +556,12 @@ def compute_matrix(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        exp_part = qml.math.exp(1j * phi)
+        arg = 1j * phi
+        if qml.math.ndim(arg) == 0:
+            return qml.math.diag(qml.math.exp(arg * np.array([0, 0, 0, 1])))
 
-        if qml.math.ndim(phi) > 0:
-            ones = qml.math.ones_like(exp_part)
-            zeros = qml.math.zeros_like(exp_part)
-            matrix = [
-                [ones, zeros, zeros, zeros],
-                [zeros, ones, zeros, zeros],
-                [zeros, zeros, ones, zeros],
-                [zeros, zeros, zeros, exp_part],
-            ]
-
-            return qml.math.stack([stack_last(row) for row in matrix], axis=-2)
-
-        return qml.math.diag([1, 1, 1, exp_part])
+        diags = qml.math.exp(qml.math.outer(arg, [0, 0, 0, 1]))
+        return diags[:, :, np.newaxis] * np.eye(4)
 
     @staticmethod
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
@@ -596,9 +593,10 @@ def compute_eigvals(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        exp_part = qml.math.exp(1j * phi)
-        ones = qml.math.ones_like(exp_part)
-        return stack_last([ones, ones, ones, exp_part])
+        if qml.math.ndim(phi) == 0:
+            return qml.math.exp(1j * phi * np.array([0, 0, 0, 1]))
+
+        return qml.math.exp(qml.math.outer(1j * phi, [0, 0, 0, 1]))
 
     @staticmethod
     def compute_decomposition(phi, wires):
@@ -870,12 +868,11 @@ def compute_matrix(theta, num_wires):  # pylint: disable=arguments-differ
             theta = qml.math.cast_like(theta, 1j)
             eigs = qml.math.cast_like(eigs, 1j)
 
-        if qml.math.ndim(theta) > 0:
-            eigvals = [qml.math.exp(-0.5j * t * eigs) for t in theta]
-            return qml.math.stack([qml.math.diag(eig) for eig in eigvals])
+        if qml.math.ndim(theta) == 0:
+            return qml.math.diag(qml.math.exp(-0.5j * theta * eigs))
 
-        eigvals = qml.math.exp(-0.5j * theta * eigs)
-        return qml.math.diag(eigvals)
+        diags = qml.math.exp(qml.math.outer(-0.5j * theta, eigs))
+        return diags[:, :, np.newaxis] * np.eye(2**num_wires)
 
     def generator(self):
         return -0.5 * functools.reduce(matmul, [qml.PauliZ(w) for w in self.wires])
@@ -915,10 +912,10 @@ def compute_eigvals(theta, num_wires):  # pylint: disable=arguments-differ
             theta = qml.math.cast_like(theta, 1j)
             eigs = qml.math.cast_like(eigs, 1j)
 
-        if qml.math.ndim(theta) > 0:
-            return qml.math.exp(qml.math.tensordot(-0.5j * theta, eigs, axes=0))
+        if qml.math.ndim(theta) == 0:
+            return qml.math.exp(-0.5j * theta * eigs)
 
-        return qml.math.exp(-0.5j * theta * eigs)
+        return qml.math.exp(qml.math.outer(-0.5j * theta, eigs))
 
     @staticmethod
     def compute_decomposition(
@@ -1662,18 +1659,12 @@ def compute_matrix(theta):  # pylint: disable=arguments-differ
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)
 
-        exp_part = qml.math.exp(-1j * theta / 2)
+        arg = 0.5j * theta
+        if qml.math.ndim(arg) == 0:
+            return qml.math.diag(qml.math.exp(arg * np.array([0, 0, -1, 1])))
 
-        ones = qml.math.ones_like(exp_part)
-        zeros = qml.math.zeros_like(exp_part)
-        matrix = [
-            [ones, zeros, zeros, zeros],
-            [zeros, ones, zeros, zeros],
-            [zeros, zeros, exp_part, zeros],
-            [zeros, zeros, zeros, qml.math.conj(exp_part)],
-        ]
-
-        return qml.math.stack([stack_last(row) for row in matrix], axis=-2)
+        diags = qml.math.exp(qml.math.outer(arg, [0, 0, -1, 1]))
+        return diags[:, :, np.newaxis] * np.eye(4)
 
     @staticmethod
     def compute_eigvals(theta):  # pylint: disable=arguments-differ
@@ -1705,10 +1696,10 @@ def compute_eigvals(theta):  # pylint: disable=arguments-differ
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)
 
-        exp_part = qml.math.exp(-0.5j * theta)
-        o = qml.math.ones_like(exp_part)
+        if qml.math.ndim(theta) == 0:
+            return qml.math.exp(0.5j * theta * np.array([0, 0, -1, 1]))
 
-        return stack_last([o, o, exp_part, qml.math.conj(exp_part)])
+        return qml.math.exp(qml.math.outer(0.5j * theta, [0, 0, -1, 1]))
 
     @staticmethod
     def compute_decomposition(phi, wires):
@@ -1996,10 +1987,12 @@ def compute_matrix(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        p = qml.math.exp(1j * phi)
-        z = qml.math.zeros_like(p)
+        arg = 1j * phi
+        if qml.math.ndim(arg) == 0:
+            return qml.math.diag(qml.math.exp(arg * np.array([0, 1])))
 
-        return qml.math.stack([stack_last([qml.math.ones_like(p), z]), stack_last([z, p])], axis=-2)
+        diags = qml.math.exp(qml.math.outer(arg, [0, 1]))
+        return diags[:, :, np.newaxis] * np.eye(2)
 
     @staticmethod
     def compute_decomposition(phi, wires):
@@ -2384,17 +2377,13 @@ def compute_matrix(phi):  # pylint: disable=arguments-differ
             s = qml.math.cast_like(s, 1j)
 
         # The following avoids casting an imaginary quantity to reals when backpropagating
-        c = (1 + 0j) * c
         js = -1j * s
-        z = qml.math.zeros_like(js)
+        if qml.math.ndim(phi) == 0:
+            return c * np.eye(4) + js * np.eye(4)[::-1]
 
-        matrix = [
-            [c, z, z, js],
-            [z, c, js, z],
-            [z, js, c, z],
-            [js, z, z, c],
-        ]
-        return qml.math.stack([stack_last(row) for row in matrix], axis=-2)
+        return qml.math.tensordot(c, np.eye(4), axes=0) + qml.math.tensordot(
+            js, np.eye(4)[::-1], axes=0
+        )
 
     @staticmethod
     def compute_decomposition(phi, wires):
@@ -2542,18 +2531,13 @@ def compute_matrix(phi):  # pylint: disable=arguments-differ
             c = qml.math.cast_like(c, 1j)
             s = qml.math.cast_like(s, 1j)
 
-        # The following avoids casting an imaginary quantity to reals when backpropagating
-        c = (1 + 0j) * c
         js = 1j * s
-        z = qml.math.zeros_like(js)
+        if qml.math.ndim(phi) == 0:
+            return c * np.eye(4) + js * np.diag([1, -1, -1, 1])[::-1]
 
-        matrix = [
-            [c, z, z, js],
-            [z, c, -js, z],
-            [z, -js, c, z],
-            [js, z, z, c],
-        ]
-        return qml.math.stack([stack_last(row) for row in matrix], axis=-2)
+        return qml.math.tensordot(c, np.eye(4), axes=0) + qml.math.tensordot(
+            js, np.diag([1, -1, -1, 1])[::-1], axes=0
+        )
 
     def adjoint(self):
         (phi,) = self.parameters
@@ -2669,18 +2653,12 @@ def compute_matrix(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        neg_phase = qml.math.exp(-0.5j * phi)
-        pos_phase = qml.math.exp(0.5j * phi)
+        arg = 0.5j * phi
+        if qml.math.ndim(arg) == 0:
+            return qml.math.diag(qml.math.exp(arg * np.array([-1, 1, 1, -1])))
 
-        zeros = qml.math.zeros_like(pos_phase)
-        matrix = [
-            [neg_phase, zeros, zeros, zeros],
-            [zeros, pos_phase, zeros, zeros],
-            [zeros, zeros, pos_phase, zeros],
-            [zeros, zeros, zeros, neg_phase],
-        ]
-
-        return qml.math.stack([stack_last(row) for row in matrix], axis=-2)
+        diags = qml.math.exp(qml.math.outer(arg, [-1, 1, 1, -1]))
+        return diags[:, :, np.newaxis] * np.eye(4)
 
     @staticmethod
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
@@ -2712,10 +2690,10 @@ def compute_eigvals(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        pos_phase = qml.math.exp(1.0j * phi / 2)
-        neg_phase = qml.math.exp(-1.0j * phi / 2)
+        if qml.math.ndim(phi) == 0:
+            return qml.math.exp(0.5j * phi * np.array([-1, 1, 1, -1]))
 
-        return stack_last([neg_phase, pos_phase, pos_phase, neg_phase])
+        return qml.math.exp(qml.math.outer(0.5j * phi, [-1, 1, 1, -1]))
 
     def adjoint(self):
         (phi,) = self.parameters
@@ -2841,14 +2819,18 @@ def compute_matrix(phi):  # pylint: disable=arguments-differ
         """
         c = qml.math.cos(phi / 2)
         s = qml.math.sin(phi / 2)
-        Y = qml.math.convert_like(np.diag([0, 1, 1, 0])[::-1].copy(), phi)
 
         if qml.math.get_interface(phi) == "tensorflow":
             c = qml.math.cast_like(c, 1j)
             s = qml.math.cast_like(s, 1j)
-            Y = qml.math.cast_like(Y, 1j)
 
-        return qml.math.diag([1, c, c, 1]) + 1j * s * Y
+        js = 1j * s
+        if qml.math.ndim(phi) == 0:
+            return qml.math.diag([1, c, c, 1]) + qml.math.diag([0, js, js, 0])[::-1]
+
+        ones = qml.math.ones_like(c)
+        diags = stack_last([ones, c, c, ones])[:, :, np.newaxis]
+        return diags * np.eye(4) + qml.math.tensordot(js, np.diag([0, 1, 1, 0])[::-1], axes=0)
 
     @staticmethod
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
@@ -2880,10 +2862,10 @@ def compute_eigvals(phi):  # pylint: disable=arguments-differ
         if qml.math.get_interface(phi) == "tensorflow":
             phi = qml.math.cast_like(phi, 1j)
 
-        pos_phase = qml.math.exp(1.0j * phi / 2)
-        neg_phase = qml.math.exp(-1.0j * phi / 2)
+        if qml.math.ndim(phi) == 0:
+            return qml.math.exp(0.5j * phi * np.array([1, -1, 0, 0]))
 
-        return qml.math.stack([pos_phase, neg_phase, 1, 1])
+        return qml.math.exp(qml.math.outer(0.5j * phi, [1, -1, 0, 0]))
 
     def adjoint(self):
         (phi,) = self.parameters