NumPy container

[tdegeus]

I would like to write my wrapper such that

1. It allocates a NumPy container.
2. Something is done the the data from a C++ function.
3. The NumPy container is returned without any more copies.

The Python API could be

import numpy as np
import mylib

a = mylib.foo(shape = [3, 4, 5])

How do I do this in the least amount of wrapper code (and with the least number of conversions)? Ideally there would be an API such that one could

.def(
    "foo",
    [](const nb::list& shape) {  
        nb::ndarray<nb::numpy, double> container(shape);
        myfunc(container.data(), container.size());
        return container;
    },
    nb::arg("shape"))

[wjakob]

You can use wrappers: module_::import_("numpy").attr("ndarray")(...) to create the initial container and use nb::cast to convert it into a nb::ndarray for manipulation. Or you can use nb::ndarray to create the initial container and call nb::cast to turn it into nb::object that can be returned (in this case, the expectation is that the nb::ndarray constructor is called with a pointer to the data, so you need to do the memory allocation yourself).

[tdegeus]

Thanks. Do you mean:

Option 1 (seems to be quite lean!)

m.def(
    "foo",
    [](const nb::list& shape) {
        auto ret = nb::cast<nb::ndarray<nb::numpy, double>>(nb::module_::import_("numpy").attr("ndarray")(shape));
        std::fill(ret.data(), ret.data() + ret.size(), 1.0);
        return ret;
    },
    nb::arg("shape"));

(Assumes #161)

Questions:

• The array appears to be zero initialised (i.e. removing the dummy code std::fill results in an array of zeros). Is that true? Can it be avoided?
• Is doing nb::module_::import_("numpy") costly when done over and over, or is it compiled away?

Option 2

Actually I don't fully get it, this seems to be what I infer from the docs, but it does not write to the memory correctly, and, frankly, look horrible

m.def(
    "foo",
    [](const std::vector<size_t>& shape) {
        size_t size = std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<size_t>{});
        size_t shape_[shape.size()];
        std::copy(shape.begin(), shape.end(), &shape_[0]);
        double *data = new double[size];
        std::fill(data, data + size, 1.0);
        nb::capsule owner(data, [](void *p) noexcept {
            delete[] (float *) p;
        });
        return nb::ndarray<nb::numpy, double>(data, shape.size(), shape_);
    },
    nb::arg("shape"));