Custom Element implementation leads to weird results

[hombit]

I was playing implementing Element for a simple case of user-defined dtype and it looks like unsafe is not really safe space for me:

use numpy::Element;
use numpy::PyArray1;
use numpy::PyArrayDescr;
use pyo3::prelude::*;

#[repr(C)]
#[derive(Clone, Copy, Debug)]
struct Record {
    x: f64,
    y: f64,
    k: i32,
}

unsafe impl Element for Record {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python<'_>) -> &PyArrayDescr {
        PyArrayDescr::new(py, &[("x", "f8"), ("y", "f8"), ("k", "i4")]).unwrap()
    }
}

#[pyfunction]
fn norm(a: &PyArray1<Record>) -> f64 {
    println!("{:?}", a);
    a.readonly().as_array().fold(0.0, |result, &record| {
        println!("result = {result}, record = {record:?}");
        result + record.x.powi(record.k) + record.y.powi(record.k)
    })
}

#[pymodule]
fn coniferest(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(norm, m)?)?;
    Ok(())
}

It is what I see when call it with Python:

result = norm(np.array([(1.0, 2.0, 2), (3.0, 4.0, 3), (5.0, 6.0, 4)],
                                       dtype=np.dtype([('x', '<f8'), ('y', '<f8'), ('k', '<i4')])))
print(f'{result = }')

array([(1., 2., 2), (3., 4., 3), (5., 6., 4)],
      dtype=[('x', '<f8'), ('y', '<f8'), ('k', '<i4')])
result = 0, record = Record { x: 1.0, y: 2.0, k: 2 }
result = 5, record = Record { x: 1.0, y: 2.0, k: 2 }
result = 10, record = Record { x: 1.0, y: 2.0, k: 2 }
result = 15.0

Do you have any idea what went wrong here?

[hombit]

If I also print a.readonly().as_array() and it looks like it has zero stride, is it supposed to be so?..

[Record { x: 1.0, y: 2.0, k: 2 }, Record { x: 1.0, y: 2.0, k: 2 }, Record { x: 1.0, y: 2.0, k: 2 }], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1

[hombit]

I'm really sorry, for some reason I missed this part of the documentation

Currently, only integer/float/complex/object types are supported

[hombit]

Just for the records: I was able to make it work with np.dtype(align=True). The full working example:

use numpy::Element;
use numpy::PyArray1;
use numpy::PyArrayDescr;
use pyo3::prelude::*;
use pyo3::py_run;
use pyo3::types::PyDict;

#[repr(C)]
#[derive(Clone, Copy, Debug)]
struct Record {
    x: f64,
    y: f64,
    k: i32,
}

fn dtype<'py>(py: Python<'py>) -> PyResult<&'py PyArrayDescr> {
    let locals = PyDict::new(py);
    py_run!(
        py,
        *locals,
        "dtype = __import__('numpy').dtype([('x', 'f8'), ('y', 'f8'), ('k', 'i4')], align=True)"
    );
    Ok(locals
        .get_item("dtype")
        .unwrap()
        .downcast::<PyArrayDescr>()?)
}

unsafe impl Element for Record {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python<'_>) -> &PyArrayDescr {
        dtype(py).unwrap()
    }
}

#[pyfunction]
fn norm(a: &PyArray1<Record>) -> f64 {
    println!("{:?}", a);
    let readonly = a.readonly();
    let array = readonly.as_array();
    println!("{:?}", array.as_standard_layout());
    array.fold(0.0, |result, record| {
        println!("result = {result}, record = {record:?}");
        result + record.x.powi(record.k) + record.y.powi(record.k)
    })
}

#[pymodule]
fn coniferest(py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(norm, m)?)?;
    m.add("dtype", dtype(py)?)?;
    Ok(())
}

[adamreichold]

See also #186