rustpyo3

Rust to Python custom type covertion with pyo3


I'm trying to create a Python package with Pyo3 & Rust.

in one of my functions, I'm trying to return a tch::Tensor type though not sure how to transform it so Python will be familiar with it.

here's what I've tried so far:

use pyo3::prelude::*;
use tch::{Tensor, Kind};
use tch::TchError;
use pyo3::exceptions::PyValueError;

struct PyTchError(TchError);

impl From<PyTchError> for PyErr {
    fn from(error: PyTchError) -> Self {
        PyValueError::new_err("error while resizing image")
    }
}

impl From<TchError> for PyTchError {
    fn from(other: TchError) -> Self {
        Self(other)
    }
}


#[pyfunction]
fn create_tensor() -> Result<Tensor, PyTchError> {
    let res = Tensor::from_slice(&[1.0, 2.0, 3.0]).to_kind(Kind::Float);
    Ok(res)
}

/// A Python module implemented in Rust.
#[pymodule]
fn pyo3_utils(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(create_tensor, m)?)?;
    Ok(())
}

which resulted in the error -

error[E0277]: the trait bound Result<tch::Tensor, PyTchError>: OkWrap<_> is not satisfied --> src/lib.rs:42:1 | 42 | #[pyfunction] | ^^^^^^^^^^^^^ the trait OkWrap<_> is not implemented for Result<tch::Tensor, PyTchError> | = help: the trait OkWrap<T> is implemented for Result<T, E>

I then tried to implement the IntoPy Trait:

struct PyTensor(Tensor);

impl IntoPy<Py<PyAny>> for PyTensor {
    fn into_py(self, py: Python<'_>) -> Py<PyAny> {
        self.0
    }
}

#[pyfunction]
fn create_tensor() -> Result<PyTensor, PyTchError> {
    let res = Tensor::from_slice(&[1.0, 2.0, 3.0]).to_kind(Kind::Float);
    Ok(PyTensor(res))
}

but still getting an error -

error[E0308]: mismatched types --> src/lib.rs:26:9 | 25 | fn into_py(self, py: Python<'_>) -> Py { |
--------- expected Py<PyAny> because of return type 26 | self.0 | ^^^^^^ expected Py<PyAny>, found Tensor |
= note: expected struct Py<PyAny> found struct tch::Tensor

error[E0277]: the trait bound Result<tch::Tensor, PyTchError>: OkWrap<_> is not satisfied --> src/lib.rs:45:1 | 45 | #[pyfunction] | ^^^^^^^^^^^^^ the trait OkWrap<_> is not implemented for Result<tch::Tensor, PyTchError> | = help: the trait OkWrap<T> is implemented for Result<T, E>

any ideas how to convert Tensor to python correctly?


Solution

  • Your IntoPy implementation doesn't make any sense: it's supposed to return a PyObject, you're returning a Tensor, the entire issue is that Tensor is not a python-compatible type.

    In this case however it doesn't really matter: tch actually provides a pyo3 bridge, so you can just depend on pyo3-tch and return a PyTensor.