Add pickle support of TwoQubitControlledUDecomposer (#14038)

* add pickle support

* add test for pickle support

* update test_assertPickle

* add impl IntoPyObject for RXXEquivalent

* add a test for pickle with dill

* update tests names

* update following review

* add back another impl

Author: ShellyGarion

crates/accelerate/src/two_qubit_decompose.rs

@@ -2479,6 +2479,28 @@ impl RXXEquivalent {
         }
     }
 }
+impl<'a, 'py> IntoPyObject<'py> for &'a RXXEquivalent {
+    type Target = PyAny;
+    type Output = Borrowed<'a, 'py, Self::Target>;
+    type Error = PyErr;
+    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
+        match self {
+            RXXEquivalent::Standard(gate) => Ok(gate.get_gate_class(py)?.bind_borrowed(py)),
+            RXXEquivalent::CustomPython(gate) => Ok(gate.as_any().bind_borrowed(py)),
+        }
+    }
+}
+impl<'py> IntoPyObject<'py> for RXXEquivalent {
+    type Target = PyAny;
+    type Output = Bound<'py, Self::Target>;
+    type Error = PyErr;
+    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
+        match self {
+            RXXEquivalent::Standard(gate) => Ok(gate.get_gate_class(py)?.bind(py).clone()),
+            RXXEquivalent::CustomPython(gate) => Ok(gate.bind(py).clone().into_any()),
+        }
+    }
+}
 
 #[derive(Clone, Debug)]
 #[pyclass(module = "qiskit._accelerate.two_qubit_decompose", subclass)]
@@ -2920,6 +2942,10 @@ impl TwoQubitControlledUDecomposer {
 
 #[pymethods]
 impl TwoQubitControlledUDecomposer {
+    fn __getnewargs__(&self) -> (&RXXEquivalent, &str) {
+        (&self.rxx_equivalent_gate, self.euler_basis.as_str())
+    }
+
     ///  Initialize the KAK decomposition.
     ///  Args:
     ///      rxx_equivalent_gate: Gate that is locally equivalent to an :class:`.RXXGate`:

test/python/synthesis/test_synthesis.py

@@ -17,6 +17,7 @@
 import contextlib
 import logging
 import math
+import dill
 import numpy as np
 import scipy
 import scipy.stats
@@ -1544,6 +1545,41 @@ def inverse(self, annotated: bool = False):
         circ = decomposer(unitary)
         self.assertEqual(Operator(unitary), Operator(circ))
 
+    def test_assert_pickle(self):
+        """Assert that TwoQubitControlledUDecomposer supports pickle"""
+
+        decomp = TwoQubitControlledUDecomposer(RXXGate)
+
+        pkl = pickle.dumps(decomp, protocol=max(4, pickle.DEFAULT_PROTOCOL))
+        decomp_cpy = pickle.loads(pkl)
+        msg_base = f"decomp:\n{decomp}\ndecomp_cpy:\n{decomp_cpy}"
+        self.assertEqual(type(decomp), type(decomp_cpy), msg_base)
+        self.assertEqual(decomp.rxx_equivalent_gate, decomp_cpy.rxx_equivalent_gate, msg=msg_base)
+        self.assertEqual(decomp.euler_basis, decomp_cpy.euler_basis, msg=msg_base)
+
+    def test_assert_pickle_with_dill(self):
+        """Assert that TwoQubitControlledUDecomposer supports pickle"""
+
+        class CustomXXGate(RXXGate):
+            """Custom RXXGate subclass that's not a standard gate"""
+
+            _standard_gate = None
+
+            def __init__(self, theta, label=None):
+                super().__init__(theta, label)
+                self.name = "MyCustomXXGate"
+
+        decomp = TwoQubitControlledUDecomposer(CustomXXGate)
+
+        pkl = dill.dumps(decomp, protocol=max(4, pickle.DEFAULT_PROTOCOL))
+        decomp_cpy = dill.loads(pkl)
+        msg_base = f"decomp:\n{decomp}\ndecomp_cpy:\n{decomp_cpy}"
+        self.assertEqual(type(decomp), type(decomp_cpy), msg_base)
+        self.assertEqual(
+            decomp.rxx_equivalent_gate.name, decomp_cpy.rxx_equivalent_gate.name, msg=msg_base
+        )
+        self.assertEqual(decomp.euler_basis, decomp_cpy.euler_basis, msg=msg_base)
+
 
 class TestDecomposeProductRaises(QiskitTestCase):
     """Check that exceptions are raised when 2q matrix is not a product of 1q unitaries"""