Returning Array with CustomType to NumPy with PyO3 + Rust-Numpy

I have defined a custom type:

#[pyclass]
#[derive(Clone, Debug)]
pub struct Dual {
    pub real: f64,
    pub vars: Arc<IndexSet<String>>,
    pub dual: Array1<f64>,
}

This struct has all the necessary operator overloads, so I can create these objects in Python, add them to NumPy arrays, and perform NumPy array operations on them, e.g. np.dot(..).

But, I have also written linalg operations in rust so I want to be able to create arrays in Python containing these types, pass them to a rust function and return to Python into a numpy array containing new Dual structs.

Calling the following function in Python causes a segmentation fault, after return:

#[pyfunction]
#[pyo3(name = "dsolve")]
fn dsolve_py<'py>(
    py: Python<'py>,
    b: PyReadonlyArray1<'py, Dual>,
) -> &'py PyArray1<Dual> {
    let b = b.as_array();  // do ops which might create a new array like c
    let c = arr1(&[Dual::new(2.0, Vec::new(), Vec::new())])
    .into_pyarray(py);
    c
}

I have added the following trait according to docs and compiler waypoints:

unsafe impl Element for Dual {
    const IS_COPY: bool = false;
    fn get_dtype<'py>(py: Python<'py>) -> &'py PyArrayDescr {
        PyArrayDescr::object(py)
    }
}

There are comments in the docs about possibly using Py<T> as the element in the array as opposed to just T, possibly also returning a PyResult<..> might be necessary. Can anyone explain in relative layman terms what is happening and going wrong?

What happens if you return Py<PyArray1<Dual>> instead of &'py PyArray1<Dual> from your function?

Immediately I get,

error[E0308]: mismatched types
  --> src/lib.rs:34:5
   |
25 | ) -> Py<PyArray1<Dual>> {
   |      ------------------ expected `pyo3::Py<PyArray<Dual, Dim<[usize; 1]>>>` because of return type
...
34 |     c
   |     ^ expected `Py<PyArray<Dual, Dim<[usize; 1]>>>`, found `&PyArray<Dual, Dim<[usize; 1]>>`

But I am trying to fix that..

I was thinking something like this:

#[pyfunction]
#[pyo3(name = "dsolve")]
fn dsolve_py<'py>(
    py: Python<'py>,
    b: PyReadonlyArray1<'py, Dual>,
) -> Py<PyArray1<Dual>> {
    let b = b.as_array();  // do ops which might create a new array like c
    let c = arr1(&[Dual::new(2.0, Vec::new(), Vec::new())])
        .into_pyarray(py);
    Py::new(py, c)
}

Yeh I tried this but Py::new(py, c)returned a PyResult so after some more expected type errors I had

... -> PyResult<Py<PyArray1<Dual>>> {
...
let c = arr1(&[Dual::new(2.0, Vec::new(), Vec::new())])
    .into_pyarray(py);
Py::new(py, c)
}

This was the closest I got but still failed becuase of a reference:

 Py::new(py, c)
   |     ^^^^^^^^^^^^^^ expected `Result<Py<PyArray<Dual, Dim<[usize; 1]>>>, PyErr>`, found `Result<Py<&PyArray<Dual, Dim<[usize; 1]>>>, PyErr>`

Hmm, tricky. This might get us nowhere, I'm sorry. I wonder if you get the segfault when you use ToPyArray instead of IntoPyArray?

#[pyfunction]
#[pyo3(name = "dsolve")]
fn dsolve_py<'py>(
    py: Python<'py>,
    b: PyReadonlyArray1<'py, Dual>,
) -> &'py PyArray1<Dual> {
    let b = b.as_array();  // do ops which might create a new array like c
    let c = arr1(&[Dual::new(2.0, Vec::new(), Vec::new())])
        .to_pyarray(py);
    c
}

This creates a copy of your array on the Python heap instead of transferring ownership of your array directly to the PyArray you return.

Yep I tried this one as well, no dice. Same problem.

I noted this page: PyArray in numpy::array - Rust in terms of Mem Location and I tried the different functions, either "Allocated by Rust" or "Allocated by NumPy", for example:

let c = arr1(&[Dual::new(2.0, Vec::new(), Vec::new())]);
let parray = PyArray::from_array(py, &c);
parray

This compiles but Seg Faults.

I wonder if this is all due to this: Element in numpy - Rust and the only implementation of the unsafe keyword that I have here:

unsafe impl Element for Dual {
    const IS_COPY: bool = false;
    fn get_dtype<'py>(py: Python<'py>) -> &'py PyArrayDescr {
        PyArrayDescr::object(py)
    }
}

In which case I would be very interested to see an example where someone can pass a custom type struct from Rust to NumPy array. I have not found an example of it.

Just found a similar report from 2021. PyArray -- make array of compound type - #6 by TimWescott

@TimWescott suggested refactoring the output to not return a PyArray but return some sequence and in Python get NumPy to convert to that to an array.

1 Like

The solution implemented was as follows.

Convert the Rust output type and return value:

#[pyfunction]
#[pyo3(name = "dsolve")]
fn dsolve_py<'py>(
    py: Python<'py>,
    b: PyReadonlyArray1<'py, Dual>,
) -> Vec<Dual> {
    let c = arr1(&[Dual::new(2.0, Vec::new(), Vec::new())]);
    c.into_raw_vec()

and on the Python side change result=dsolve(b) to result=np.array(dsolve(b))

1 Like

According to Element's documentation you can convert a ndarray::Array<Py<YourType>, D> into a &PyArray<PyObject> using PyArray::from_owned_object_array, that might allow you to avoid the explicit call to np.array in python.

I saw this too but I am not very skilled in Rust and it wasnt immediately obvious how to convert an entire linalg module that is based on Array<Dual> types into Array<Py<Dual>> dtypes.
Perhaps I need to implement some kind of converter. But then this is the same kind of monkey-patching as below.

I viewed it as more work and more confusing for my Rust code than to just provide a monkey-patching solution to pass List[Dual] and Vec<Dual> between Python and Rust, and either use numpy.to_list() and numpy.from_list(), on the Python side. Im sure I lose efficiency with this but this efficiency gains of rust in the first place way outstrip this loss.

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.