Skip to content

Commit

Permalink
Use pickle __reduce__ for TwoQubitWeylDecomposition (#12003)
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@jakelishman
jakelishman committed Mar 16, 2024
1 parent af659e2 commit 6e280c1
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 122 deletions.
49 changes: 28 additions & 21 deletions crates/accelerate/src/euler_one_qubit_decomposer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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<Self> {
#[pymethods]
impl EulerBasis {
fn __reduce__(&self, py: Python) -> Py<PyAny> {
(py.get_type::<Self>(), (PyString::new(py, self.as_str()),)).into_py(py)
}

#[new]
pub fn from_str(input: &str) -> PyResult<Self> {
let res = match input {
"U321" => EulerBasis::U321,
"U3" => EulerBasis::U3,
Expand All @@ -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}'"
)));
}
};
Expand Down Expand Up @@ -702,7 +708,7 @@ pub fn unitary_to_gate_sequence(
) -> PyResult<Option<OneQubitGateSequence>> {
let mut target_basis_vec: Vec<EulerBasis> = 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();
Expand Down Expand Up @@ -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::<OneQubitGateSequence>()?;
m.add_class::<OneQubitGateErrorMap>()?;
m.add_class::<EulerBasis>()?;
Ok(())
}
177 changes: 77 additions & 100 deletions crates/accelerate/src/two_qubit_decompose.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<Py<PyAny>> {
// 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::<Self>().getattr("_from_u8")?, (val,)).into_py(py))
}

fn from_u8(val: u8) -> Self {
#[staticmethod]
fn _from_u8(val: u8) -> PyResult<Self> {
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}'"
))),
}
}
}
Expand Down Expand Up @@ -470,91 +477,61 @@ const IPX: [[Complex64; 2]; 2] = [

#[pymethods]
impl TwoQubitWeylDecomposition {
fn __getstate__(&self, py: Python) -> ([f64; 5], [PyObject; 5], u8, String, Option<f64>) {
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<Complex64>; 5],
specialization: Specialization,
default_euler_basis: EulerBasis,
calculated_fidelity: f64,
requested_fidelity: Option<f64>,
) -> 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<Complex64>; 5],
u8,
String,
Option<f64>,
),
) {
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<f64>, Option<Specialization>, bool) {
(
self.unitary_matrix.to_pyarray(py).into(),
self.requested_fidelity,
None,
true,
fn __reduce__(&self, py: Python) -> PyResult<Py<PyAny>> {
Ok((
py.get_type::<Self>().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<Complex64>,
fidelity: Option<f64>,
_specialization: Option<Specialization>,
_pickle_context: bool,
) -> PyResult<Self> {
// 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<Complex64> = aview2(&IPZ);
let ipy: ArrayView2<Complex64> = aview2(&IPY);
let ipx: ArrayView2<Complex64> = aview2(&IPX);
Expand Down Expand Up @@ -1066,7 +1043,7 @@ impl TwoQubitWeylDecomposition {
atol: Option<f64>,
) -> PyResult<TwoQubitGateSequence> {
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<EulerBasis> = vec![euler_basis];
Expand Down
2 changes: 1 addition & 1 deletion qiskit/synthesis/two_qubit/two_qubit_decompose.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,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`."""
Expand Down

0 comments on commit 6e280c1

Please sign in to comment.