Pennylane-Qualtran interoperability (#1559)

* bloq as op

* basic gates (most of them)

* identity gate

* rotation gates

* address comments

* some more gates

* fix typing

* change imports around

* import paths change

* undo

* import paths change

* import paths change

* tests

* formatting

* fix mypy, pylint, etc

* extra space

* fix

* fix typo in the documentation of CModMulK (#1594)

* trigger ci

* fix back

* fix

* add tests for matrix comparisons

---------

Co-authored-by: Noureldin <noureldinyosri@google.com>

Author: austingmhuang

dev_tools/conf/mypy.ini

@@ -21,7 +21,7 @@ follow_imports = silent
 ignore_missing_imports = true
 
 # Non-Google
-[mypy-sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,plotly.*,dash.*,tensorflow_docs.*,fxpmath.*,ipywidgets.*,cachetools.*,pydot.*,nbformat.*,nbconvert.*,openfermion.*]
+[mypy-sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,plotly.*,dash.*,tensorflow_docs.*,fxpmath.*,ipywidgets.*,cachetools.*,pydot.*,nbformat.*,nbconvert.*,openfermion.*,pennylane.*]
 follow_imports = silent
 ignore_missing_imports = true
 

qualtran/_infra/bloq.py

@@ -21,9 +21,12 @@
 if TYPE_CHECKING:
     import cirq
     import networkx as nx
+    import pennylane as qml
     import quimb.tensor as qtn
     import sympy
     from numpy.typing import NDArray
+    from pennylane.operation import Operation
+    from pennylane.wires import Wires
 
     from qualtran import (
         AddControlledT,
@@ -466,6 +469,25 @@ def as_cirq_op(
             bloq=self, cirq_quregs=cirq_quregs, qubit_manager=qubit_manager
         )
 
+    def as_pl_op(self, wires: 'Wires') -> 'Operation':
+        """Override this method to support conversion to a PennyLane operation.
+
+        If this method is not overriden, the default implementation will wrap this bloq
+        in a `FromBloq` shim.
+
+        Args:
+            wires: the wires that the op acts on
+
+        Returns:
+            ~.Operation: A PennyLane operation corresponding to this bloq acting on the
+                provided wires or None. This method should return None if and only if the bloq
+                instance truly should not be included in the PennyLane circuit (e.g. for reshaping
+                bloqs). A bloq with no PennyLane equivalent should raise an exception instead.
+        """
+        from pennylane.io import FromBloq
+
+        return FromBloq(bloq=self, wires=wires)
+
     def on(self, *qubits: 'cirq.Qid') -> 'cirq.Operation':
         """A `cirq.Operation` of this bloq operating on the given qubits.
 

qualtran/_infra/bloq_test.py

@@ -34,6 +34,14 @@ def test_bloq():
         tb.decompose_bloq()
 
 
+def test_as_pl_op():
+    import pennylane as qml
+
+    tb = TestTwoBitOp()
+
+    assert tb.as_pl_op(wires=[0, 1]) == qml.FromBloq(TestTwoBitOp(), wires=[0, 1])
+
+
 def test_as_composite_bloq():
     tb = TestAtom()
     with pytest.raises(DecomposeTypeError):

qualtran/bloqs/basic_gates/cnot.py

@@ -38,6 +38,8 @@
 if TYPE_CHECKING:
     import cirq
     import quimb.tensor as qtn
+    from pennylane.operation import Operation
+    from pennylane.wires import Wires
 
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.simulation.classical_sim import ClassicalValT
@@ -122,6 +124,11 @@ def as_cirq_op(
         (target,) = target
         return cirq.CNOT(ctrl, target), {'ctrl': np.array([ctrl]), 'target': np.array([target])}
 
+    def as_pl_op(self, wires: 'Wires') -> 'Operation':
+        import pennylane as qml
+
+        return qml.CNOT(wires=wires)
+
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None:
             return Text('')

qualtran/bloqs/basic_gates/cnot_test.py

@@ -14,6 +14,7 @@
 
 import cirq
 import numpy as np
+import pennylane as qml
 import pytest
 
 from qualtran import BloqBuilder, Signature
@@ -39,6 +40,19 @@ def test_cnot_tensor():
     np.testing.assert_allclose(should_be, matrix)
 
 
+def test_cnot_vs_pl():
+    bloq = CNOT()
+    matrix = qml.FromBloq(bloq, wires=[0, 1]).matrix()
+    # fmt: off
+    should_be = np.array([
+        [1, 0, 0, 0],
+        [0, 1, 0, 0],
+        [0, 0, 0, 1],
+        [0, 0, 1, 0]])
+    # fmt: on
+    np.testing.assert_allclose(should_be, matrix)
+
+
 def test_cnot_cbloq_tensor_vs_cirq():
     bb, soqs = BloqBuilder.from_signature(Signature.build(c=1, t=1))
     c, t = bb.add(CNOT(), ctrl=soqs['c'], target=soqs['t'])

qualtran/bloqs/basic_gates/global_phase.py

@@ -35,6 +35,8 @@
 
 if TYPE_CHECKING:
     import quimb.tensor as qtn
+    from pennylane.operation import Operation
+    from pennylane.wires import Wires
 
 
 @frozen(kw_only=True)
@@ -78,6 +80,12 @@ def from_coefficient(
     def cirq_gate(self) -> cirq.Gate:
         return cirq.GlobalPhaseGate(self.coefficient)
 
+    def as_pl_op(self, wires: 'Wires') -> 'Operation':
+        import numpy as np
+        import pennylane as qml
+
+        return qml.GlobalPhase(phi=self.exponent * np.pi, wires=wires)
+
     def decompose_bloq(self) -> 'CompositeBloq':
         raise DecomposeTypeError(f"{self} is atomic")
 

qualtran/bloqs/basic_gates/global_phase_test.py

@@ -14,6 +14,7 @@
 
 import cirq
 import numpy as np
+import pennylane as qml
 import pytest
 
 from qualtran import CtrlSpec
@@ -54,6 +55,22 @@ def test_cirq_interop():
     assert cirq_gate_to_bloq(gate) == bloq
 
 
+def test_pl_interop():
+    bloq = GlobalPhase(exponent=0.5)
+    pl_op_from_bloq = bloq.as_pl_op(wires=[0])
+    pl_op = qml.GlobalPhase(phi=0.5 * np.pi, wires=[0])
+    assert pl_op_from_bloq == pl_op
+
+    matrix = pl_op_from_bloq.matrix()
+    # fmt: off
+    should_be = np.array([
+        [6.123234e-17-1.j, 0.000000e+00+0.j],
+        [0.000000e+00+0.j, 6.123234e-17-1.j]
+    ])
+    # fmt: on
+    np.testing.assert_allclose(should_be, matrix)
+
+
 def test_t_complexity():
     assert GlobalPhase(exponent=0.5).t_complexity() == TComplexity()
 

qualtran/bloqs/basic_gates/hadamard.py

@@ -37,6 +37,8 @@
 if TYPE_CHECKING:
     import cirq
     import quimb.tensor as qtn
+    from pennylane.operation import Operation
+    from pennylane.wires import Wires
 
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.resource_counting import CostKey
@@ -105,6 +107,11 @@ def as_cirq_op(
         (q,) = q
         return cirq.H(q), {'q': np.array([q])}
 
+    def as_pl_op(self, wires: 'Wires') -> 'Operation':
+        import pennylane as qml
+
+        return qml.Hadamard(wires=wires)
+
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None:
             return Text('')

qualtran/bloqs/basic_gates/hadamard_test.py

@@ -70,6 +70,19 @@ def test_cirq_interop():
     assert cirq_gate_to_bloq(op.gate) == CHadamard()
 
 
+def test_pl_interop():
+    import pennylane as qml
+
+    bloq = Hadamard()
+    pl_op_from_bloq = bloq.as_pl_op(wires=[0])
+    pl_op = qml.Hadamard(wires=[0])
+    assert pl_op_from_bloq == pl_op
+
+    matrix = pl_op.matrix()
+    should_be = bloq.tensor_contract()
+    np.testing.assert_allclose(should_be, matrix)
+
+
 def test_active_chadamard_is_hadamard():
     bb = BloqBuilder()
     q = bb.add_register('q', 1)

qualtran/bloqs/basic_gates/identity.py

@@ -38,6 +38,8 @@
 if TYPE_CHECKING:
     import cirq
     import quimb.tensor as qtn
+    from pennylane.operation import Operation
+    from pennylane.wires import Wires
 
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.simulation.classical_sim import ClassicalValT
@@ -88,6 +90,11 @@ def as_cirq_op(
 
         return cirq.IdentityGate(self.bitsize).on(*q), {'q': q}
 
+    def as_pl_op(self, wires: 'Wires') -> 'Operation':
+        import pennylane as qml
+
+        return qml.Identity(wires=wires)
+
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None:
             return Text('')