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

help
1

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?

2

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

3

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..

4

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)
}
5

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>`
6

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.

7

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.

8

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.

9

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))

10

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.

11

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.