allow creating a prod op with a qfunc (#4011)

* allow creating a prod op with a qfunc

* changelog

* prod(qfunc) now returns a callable that accepts args

* update docstring example

* update prod wrapper to return operand itself if only one provided

* update and test typeerror when prod arg is not callable

doc/releases/changelog-dev.md

@@ -18,6 +18,9 @@
 * `qml.specs` is compatible with custom operations that have `depth` bigger than 1.
   [(#4033)](https://github.com/PennyLaneAI/pennylane/pull/4033)
 
+* `qml.prod` now accepts a single qfunc input for creating new `Prod` operators.
+  [(#4011)](https://github.com/PennyLaneAI/pennylane/pull/4011)
+
 <h3>Breaking changes 💔</h3>
 
 * `pennylane.collections`, `pennylane.op_sum`, and `pennylane.utils.sparse_hamiltonian` are removed.
@@ -40,7 +43,6 @@ Isaac De Vlugt,
 Soran Jahangiri,
 Christina Lee,
 Mudit Pandey,
-Christina Lee,
-Mudit Pandey,
+Matthew Silverman,
 Jay Soni
 

pennylane/ops/op_math/prod.py

@@ -17,11 +17,10 @@
 """
 import itertools
 from copy import copy
-from functools import reduce
+from functools import reduce, wraps
 from itertools import combinations
 from typing import List, Tuple, Union
 
-import numpy as np
 from scipy.sparse import kron as sparse_kron
 
 import pennylane as qml
@@ -33,6 +32,7 @@
 from pennylane.ops.qubit import Hamiltonian
 from pennylane.ops.qubit.non_parametric_ops import PauliX, PauliY, PauliZ
 from pennylane.queuing import QueuingManager
+from pennylane.typing import TensorLike
 from pennylane.wires import Wires
 
 from .composite import CompositeOp
@@ -51,7 +51,8 @@ def prod(*ops, do_queue=True, id=None, lazy=True):
     that the given operators act on.
 
     Args:
-        ops (tuple[~.operation.Operator]): The operators we would like to multiply
+        ops (Union[tuple[~.operation.Operator], Callable]): The operators we would like to multiply.
+            Alternatively, a single qfunc that queues operators can be passed to this function.
 
     Keyword Args:
         do_queue (bool): determines if the product operator will be queued. Default is True.
@@ -84,7 +85,32 @@ def prod(*ops, do_queue=True, id=None, lazy=True):
     >>> prod_op.matrix()
     array([[ 0, -1],
            [ 1,  0]])
+
+    You can also create a prod operator by passing a qfunc to prod, like the following:
+
+    >>> def qfunc(x):
+    ...     qml.RX(x, 0)
+    ...     qml.CNOT([0, 1])
+    >>> prod_op = prod(qfunc)(1.1)
+    >>> prod_op
+    CNOT(wires=[0, 1]) @ RX(1.1, wires=[0])
     """
+    if len(ops) == 1:
+        if isinstance(ops[0], qml.operation.Operator):
+            return ops[0]
+
+        fn = ops[0]
+
+        if not callable(fn):
+            raise TypeError(f"Unexpected argument of type {type(fn).__name__} passed to qml.prod")
+
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            qs = qml.tape.make_qscript(fn)(*args, **kwargs)
+            return prod(*qs.operations[::-1], do_queue=do_queue, id=id, lazy=lazy)
+
+        return wrapper
+
     if lazy:
         return Prod(*ops, do_queue=do_queue, id=id)
 
@@ -257,7 +283,7 @@ def decomposition(self):
     def matrix(self, wire_order=None):
         """Representation of the operator as a matrix in the computational basis."""
 
-        mats: List[np.ndarray] = []  # TODO: change type to `tensor_like` when available
+        mats: List[TensorLike] = []
         batched: List[bool] = []  # batched[i] tells if mats[i] is batched or not
         for ops in self.overlapping_ops:
             gen = (

tests/ops/op_math/test_prod.py

@@ -116,7 +116,7 @@ class MyOp(qml.RX):  # pylint:disable=too-few-public-methods
     has_diagonalizing_gates = False
 
 
-class TestInitialization:
+class TestInitialization:  # pylint:disable=too-many-public-methods
     """Test the initialization."""
 
     @pytest.mark.parametrize("id", ("foo", "bar"))
@@ -316,6 +316,96 @@ def test_has_diagonalizing_gates_false_via_factor(self):
         prod_op = prod(MyOp(3.1, 0), qml.PauliX(2), do_queue=True)
         assert prod_op.has_diagonalizing_gates is False
 
+    def test_qfunc_init(self):
+        """Tests prod initialization with a qfunc argument."""
+
+        def qfunc():
+            qml.Hadamard(0)
+            qml.CNOT([0, 1])
+            qml.RZ(1.1, 1)
+
+        prod_gen = prod(qfunc)
+        assert callable(prod_gen)
+        prod_op = prod_gen()
+        expected = prod(qml.RZ(1.1, 1), qml.CNOT([0, 1]), qml.Hadamard(0))
+        assert qml.equal(prod_op, expected)
+        assert prod_op.wires == Wires([1, 0])
+
+    def test_qfunc_init_accepts_args_kwargs(self):
+        """Tests that prod preserves args when wrapping qfuncs."""
+
+        def qfunc(x, run_had=False):
+            if run_had:
+                qml.Hadamard(0)
+            qml.RX(x, 0)
+            qml.CNOT([0, 1])
+
+        prod_gen = prod(qfunc)
+        assert qml.equal(prod_gen(1.1), prod(qml.CNOT([0, 1]), qml.RX(1.1, 0)))
+        assert qml.equal(
+            prod_gen(2.2, run_had=True), prod(qml.CNOT([0, 1]), qml.RX(2.2, 0), qml.Hadamard(0))
+        )
+
+    def test_qfunc_init_propagates_Prod_kwargs(self):
+        """Tests that additional kwargs for Prod are propagated using qfunc initialization."""
+
+        def qfunc(x):
+            qml.prod(qml.RX(x, 0), qml.PauliZ(1))
+            qml.CNOT([0, 1])
+
+        prod_gen = prod(qfunc, do_queue=False, id=123987, lazy=False)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            prod_op = prod_gen(1.1)
+
+        assert prod_op not in q  # do_queue worked
+        assert prod_op.id == 123987  # id was set
+        assert qml.equal(prod_op, prod(qml.CNOT([0, 1]), qml.PauliZ(1), qml.RX(1.1, 0)))  # eager
+
+    def test_qfunc_init_only_works_with_one_qfunc(self):
+        """Test that the qfunc init only occurs when one callable is passed to prod."""
+
+        def qfunc():
+            qml.Hadamard(0)
+            qml.CNOT([0, 1])
+
+        prod_op = prod(qfunc)()
+        assert qml.equal(prod_op, prod(qml.CNOT([0, 1]), qml.Hadamard(0)))
+
+        def fn2():
+            qml.PauliX(0)
+            qml.PauliY(1)
+
+        for args in [(qfunc, fn2), (qfunc, qml.PauliX), (qml.PauliX, qfunc)]:
+            with pytest.raises(AttributeError, match="has no attribute 'wires'"):
+                prod(*args)
+
+    def test_qfunc_init_returns_single_op(self):
+        """Tests that if a qfunc only queues one operator, that operator is returned."""
+
+        def qfunc():
+            qml.PauliX(0)
+
+        prod_op = prod(qfunc)()
+        assert qml.equal(prod_op, qml.PauliX(0))
+        assert not isinstance(prod_op, Prod)
+
+    def test_prod_accepts_single_operator_but_Prod_does_not(self):
+        """Tests that the prod wrapper can accept a single operator, and return it."""
+
+        x = qml.PauliX(0)
+        prod_op = prod(x)
+        assert prod_op is x
+        assert not isinstance(prod_op, Prod)
+
+        with pytest.raises(ValueError, match="Require at least two operators"):
+            Prod(x)
+
+    def test_prod_fails_with_non_callable_arg(self):
+        """Tests that prod explicitly checks that a single-arg is either an Operator or callable."""
+        with pytest.raises(TypeError, match="Unexpected argument of type int passed to qml.prod"):
+            prod(1)
+
 
 class TestMatrix:
     """Test matrix-related methods."""