From d86b391ce2c3463783e73a6ba57079ca530bbfeb Mon Sep 17 00:00:00 2001 From: Jake Lishman Date: Wed, 13 Mar 2024 18:39:39 +0000 Subject: [PATCH] Use pickle `__reduce__` for `TwoQubitWeylDecomposition` Since this class does complicated calculation within its `__new__` method (as many extension classes do), we want to use an alternative extension-module constructor for the object on unpickle, to avoid leaking the "pickle" / "no pickle" state through the default constructor. This implements a private Python-space method that directly constructs the object from components, which is then used as the `__reduce__` target, along with the pickle state. --- .../src/euler_one_qubit_decomposer.rs | 49 ++--- crates/accelerate/src/two_qubit_decompose.rs | 177 ++++++++---------- .../two_qubit/two_qubit_decompose.py | 2 +- 3 files changed, 106 insertions(+), 122 deletions(-) diff --git a/crates/accelerate/src/euler_one_qubit_decomposer.rs b/crates/accelerate/src/euler_one_qubit_decomposer.rs index 6af60400a910..5f9f479f6006 100644 --- a/crates/accelerate/src/euler_one_qubit_decomposer.rs +++ b/crates/accelerate/src/euler_one_qubit_decomposer.rs @@ -19,8 +19,9 @@ use smallvec::{smallvec, SmallVec}; use std::cmp::Ordering; use std::f64::consts::PI; -use pyo3::exceptions::{PyIndexError, PyTypeError}; +use pyo3::exceptions::{PyIndexError, PyValueError}; use pyo3::prelude::*; +use pyo3::types::PyString; use pyo3::wrap_pyfunction; use pyo3::Python; @@ -571,28 +572,33 @@ pub enum EulerBasis { XZX, } -#[pymethods] impl EulerBasis { - #![allow(clippy::wrong_self_convention)] - pub fn to_str(&self) -> String { + pub fn as_str(&self) -> &'static str { match self { - Self::U321 => "U321".to_string(), - Self::U3 => "U3".to_string(), - Self::U => "U".to_string(), - Self::PSX => "PSX".to_string(), - Self::ZSX => "ZSX".to_string(), - Self::ZSXX => "ZSXX".to_string(), - Self::U1X => "U1X".to_string(), - Self::RR => "RR".to_string(), - Self::ZYZ => "ZYZ".to_string(), - Self::ZXZ => "ZXZ".to_string(), - Self::XYX => "XYX".to_string(), - Self::XZX => "XZX".to_string(), + Self::U321 => "U321", + Self::U3 => "U3", + Self::U => "U", + Self::PSX => "PSX", + Self::ZSX => "ZSX", + Self::ZSXX => "ZSXX", + Self::U1X => "U1X", + Self::RR => "RR", + Self::ZYZ => "ZYZ", + Self::ZXZ => "ZXZ", + Self::XYX => "XYX", + Self::XZX => "XZX", } } +} - #[staticmethod] - pub fn from_string(input: &str) -> PyResult { +#[pymethods] +impl EulerBasis { + fn __reduce__(&self, py: Python) -> Py { + (py.get_type::(), (PyString::new(py, self.as_str()),)).into_py(py) + } + + #[new] + pub fn from_str(input: &str) -> PyResult { let res = match input { "U321" => EulerBasis::U321, "U3" => EulerBasis::U3, @@ -607,8 +613,8 @@ impl EulerBasis { "XYX" => EulerBasis::XYX, "XZX" => EulerBasis::XZX, basis => { - return Err(PyTypeError::new_err(format!( - "Invalid target basis {basis}" + return Err(PyValueError::new_err(format!( + "Invalid target basis '{basis}'" ))); } }; @@ -702,7 +708,7 @@ pub fn unitary_to_gate_sequence( ) -> PyResult> { let mut target_basis_vec: Vec = Vec::with_capacity(target_basis_list.len()); for basis in target_basis_list { - let basis_enum = EulerBasis::from_string(basis)?; + let basis_enum = EulerBasis::from_str(basis)?; target_basis_vec.push(basis_enum) } let unitary_mat = unitary.as_array(); @@ -880,5 +886,6 @@ pub fn euler_one_qubit_decomposer(_py: Python, m: &PyModule) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(compute_error_list))?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/crates/accelerate/src/two_qubit_decompose.rs b/crates/accelerate/src/two_qubit_decompose.rs index cb3c1308829e..6c4744174726 100644 --- a/crates/accelerate/src/two_qubit_decompose.rs +++ b/crates/accelerate/src/two_qubit_decompose.rs @@ -20,7 +20,7 @@ use approx::abs_diff_eq; use num_complex::{Complex, Complex64, ComplexFloat}; -use pyo3::exceptions::PyIndexError; +use pyo3::exceptions::{PyIndexError, PyValueError}; use pyo3::import_exception; use pyo3::prelude::*; use pyo3::wrap_pyfunction; @@ -359,35 +359,42 @@ enum Specialization { fSimabmbEquiv, } +#[pymethods] impl Specialization { - fn to_u8(self) -> u8 { - match self { - Specialization::General => 0, - Specialization::IdEquiv => 1, - Specialization::SWAPEquiv => 2, - Specialization::PartialSWAPEquiv => 3, - Specialization::PartialSWAPFlipEquiv => 4, - Specialization::ControlledEquiv => 5, - Specialization::MirrorControlledEquiv => 6, - Specialization::fSimaabEquiv => 7, - Specialization::fSimabbEquiv => 8, - Specialization::fSimabmbEquiv => 9, - } + fn __reduce__(&self, py: Python) -> PyResult> { + // Ideally we'd use the string-only form of `__reduce__` for simplicity, but PyO3 enums + // don't produce Python singletons, and pickle doesn't like that. + let val: u8 = match self { + Self::General => 0, + Self::IdEquiv => 1, + Self::SWAPEquiv => 2, + Self::PartialSWAPEquiv => 3, + Self::PartialSWAPFlipEquiv => 4, + Self::ControlledEquiv => 5, + Self::MirrorControlledEquiv => 6, + Self::fSimaabEquiv => 7, + Self::fSimabbEquiv => 8, + Self::fSimabmbEquiv => 9, + }; + Ok((py.get_type::().getattr("_from_u8")?, (val,)).into_py(py)) } - fn from_u8(val: u8) -> Self { + #[staticmethod] + fn _from_u8(val: u8) -> PyResult { match val { - 0 => Specialization::General, - 1 => Specialization::IdEquiv, - 2 => Specialization::SWAPEquiv, - 3 => Specialization::PartialSWAPEquiv, - 4 => Specialization::PartialSWAPFlipEquiv, - 5 => Specialization::ControlledEquiv, - 6 => Specialization::MirrorControlledEquiv, - 7 => Specialization::fSimaabEquiv, - 8 => Specialization::fSimabbEquiv, - 9 => Specialization::fSimabmbEquiv, - _ => unreachable!("Invalid specialization value"), + 0 => Ok(Self::General), + 1 => Ok(Self::IdEquiv), + 2 => Ok(Self::SWAPEquiv), + 3 => Ok(Self::PartialSWAPEquiv), + 4 => Ok(Self::PartialSWAPFlipEquiv), + 5 => Ok(Self::ControlledEquiv), + 6 => Ok(Self::MirrorControlledEquiv), + 7 => Ok(Self::fSimaabEquiv), + 8 => Ok(Self::fSimabbEquiv), + 9 => Ok(Self::fSimabmbEquiv), + x => Err(PyValueError::new_err(format!( + "unknown specialization discriminant '{x}'" + ))), } } } @@ -470,91 +477,61 @@ const IPX: [[Complex64; 2]; 2] = [ #[pymethods] impl TwoQubitWeylDecomposition { - fn __getstate__(&self, py: Python) -> ([f64; 5], [PyObject; 5], u8, String, Option) { - let specialization = self.specialization.to_u8(); - ( - [ - self.a, - self.b, - self.c, - self.global_phase, - self.calculated_fidelity, - ], - [ - self.K1l.to_pyarray(py).into(), - self.K1r.to_pyarray(py).into(), - self.K2l.to_pyarray(py).into(), - self.K2r.to_pyarray(py).into(), - self.unitary_matrix.to_pyarray(py).into(), - ], + #[staticmethod] + fn _from_state( + angles: [f64; 4], + matrices: [PyReadonlyArray2; 5], + specialization: Specialization, + default_euler_basis: EulerBasis, + calculated_fidelity: f64, + requested_fidelity: Option, + ) -> Self { + let [a, b, c, global_phase] = angles; + Self { + a, + b, + c, + global_phase, + K1l: matrices[0].as_array().to_owned(), + K1r: matrices[1].as_array().to_owned(), + K2l: matrices[2].as_array().to_owned(), + K2r: matrices[3].as_array().to_owned(), specialization, - self.default_euler_basis.to_str(), - self.requested_fidelity, - ) - } - - fn __setstate__( - &mut self, - state: ( - [f64; 5], - [PyReadonlyArray2; 5], - u8, - String, - Option, - ), - ) { - self.a = state.0[0]; - self.b = state.0[1]; - self.c = state.0[2]; - self.global_phase = state.0[3]; - self.calculated_fidelity = state.0[4]; - self.K1l = state.1[0].as_array().to_owned(); - self.K1r = state.1[1].as_array().to_owned(); - self.K2l = state.1[2].as_array().to_owned(); - self.K2r = state.1[3].as_array().to_owned(); - self.unitary_matrix = state.1[4].as_array().to_owned(); - self.default_euler_basis = EulerBasis::from_string(&state.3).unwrap(); - self.requested_fidelity = state.4; - self.specialization = Specialization::from_u8(state.2); + default_euler_basis, + calculated_fidelity, + requested_fidelity, + unitary_matrix: matrices[4].as_array().to_owned(), + } } - fn __getnewargs__(&self, py: Python) -> (PyObject, Option, Option, bool) { - ( - self.unitary_matrix.to_pyarray(py).into(), - self.requested_fidelity, - None, - true, + fn __reduce__(&self, py: Python) -> PyResult> { + Ok(( + py.get_type::().getattr("_from_state")?, + ( + [self.a, self.b, self.c, self.global_phase], + [ + self.K1l.to_pyarray(py), + self.K1r.to_pyarray(py), + self.K2l.to_pyarray(py), + self.K2r.to_pyarray(py), + self.unitary_matrix.to_pyarray(py), + ], + self.specialization, + self.default_euler_basis, + self.calculated_fidelity, + self.requested_fidelity, + ), ) + .into_py(py)) } #[new] - #[pyo3(signature=(unitary_matrix, fidelity=DEFAULT_FIDELITY, _specialization=None, _pickle_context=false))] + #[pyo3(signature=(unitary_matrix, fidelity=DEFAULT_FIDELITY, _specialization=None))] fn new( unitary_matrix: PyReadonlyArray2, fidelity: Option, _specialization: Option, - _pickle_context: bool, ) -> PyResult { - // If we're in a pickle context just make the closest to an empty - // object as we can with minimal allocations and effort. All the - // data will be filled in during deserialization from __setstate__. - if _pickle_context { - return Ok(TwoQubitWeylDecomposition { - a: 0., - b: 0., - c: 0., - global_phase: 0., - K1l: Array2::zeros((0, 0)), - K2l: Array2::zeros((0, 0)), - K1r: Array2::zeros((0, 0)), - K2r: Array2::zeros((0, 0)), - specialization: Specialization::General, - default_euler_basis: EulerBasis::ZYZ, - requested_fidelity: fidelity, - calculated_fidelity: 0., - unitary_matrix: Array2::zeros((0, 0)), - }); - } let ipz: ArrayView2 = aview2(&IPZ); let ipy: ArrayView2 = aview2(&IPY); let ipx: ArrayView2 = aview2(&IPX); @@ -1066,7 +1043,7 @@ impl TwoQubitWeylDecomposition { atol: Option, ) -> PyResult { let euler_basis: EulerBasis = match euler_basis { - Some(basis) => EulerBasis::from_string(basis)?, + Some(basis) => EulerBasis::from_str(basis)?, None => self.default_euler_basis, }; let target_1q_basis_list: Vec = vec![euler_basis]; diff --git a/qiskit/synthesis/two_qubit/two_qubit_decompose.py b/qiskit/synthesis/two_qubit/two_qubit_decompose.py index d8d6b0d8033d..3467cc33ec08 100644 --- a/qiskit/synthesis/two_qubit/two_qubit_decompose.py +++ b/qiskit/synthesis/two_qubit/two_qubit_decompose.py @@ -248,7 +248,7 @@ def from_bytes( bytes_in: bytes, *, requested_fidelity: float, - _specialization: two_qubit_decompose.Specialization | None, + _specialization: two_qubit_decompose.Specialization | None = None, **kwargs, ) -> "TwoQubitWeylDecomposition": """Decode bytes into :class:`.TwoQubitWeylDecomposition`."""