Return nb::ndarray that differ by dtype

[senyai]

This does not compile.

  m.def("create_2d", [](size_t rows, size_t cols, bool integral) {
    if (integral) {
      int *data = new int[rows * cols];
      for (size_t i = 0; i < rows * cols; ++i) data[i] = (int)i;

      // Delete 'data' when the 'owner' capsule expires
      nb::capsule owner(data, [](void *p) noexcept { delete[] (int *)p; });

      return nb::ndarray<nb::numpy, int, nb::ndim<2>>(
          /* data = */ data,
          /* shape = */ {rows, cols},
          /* owner = */ owner);
    } else {
      float *data = new float[rows * cols];
      for (size_t i = 0; i < rows * cols; ++i) data[i] = (float)i;

      // Delete 'data' when the 'owner' capsule expires
      nb::capsule owner(data, [](void *p) noexcept { delete[] (float *)p; });

      return nb::ndarray<nb::numpy, float, nb::ndim<2>>(
          /* data = */ data,
          /* shape = */ {rows, cols},
          /* owner = */ owner);
    }
  });

error: inconsistent types ‘nanobind::ndarray<nanobind::numpy, int, nanobind::detail::shape<-1, -1> >’ and ‘nanobind::ndarray<nanobind::numpy, float, nanobind::detail::shape<-1, -1> >’ deduced for lambda return type

[hpkfft]

That's a C++ error that is unrelated to nanobind. The following does not compile:

    auto lambda = [](std::size_t rows, bool integral) {
            if (integral) {
                return 1;
            } else {
                return 2.0;
            }
        };

with a similar error message:

error: inconsistent types 'int' and 'double' deduced for lambda return type

[senyai]

I know how to change your example so it compiles:

...
if (integral) {
  return PyLong_FromLong(1);
} else {
  return PyFloat_FromDouble(2.0);
}

Now I need something similar for my example.

[hpkfft]

I cannot guess what you want. A lambda has to return the same type from all branches the code takes.
Perhaps make your common type an ndarray of bytes.
Or, since you're willing to use the Python C API directly, don't use ndarray at all.

Typically, the best solution would be to use templates, writing

template <typename T>
nb::ndarray<nb::numpy, T, nb::ndim<2>> create_2d(std::size_t rows, std::size_t cols);

and binding that for both int and float.

Or, make a std::variant of the two ndarray types and return that.

[hpkfft]

Oh, maybe you want to return an nb::object? If so, use nb::cast.
https://nanobind.readthedocs.io/en/latest/api_core.html#casting
return nb::cast(nb::ndarray<nb::numpy, float, nb::ndim<2>>(.....))

[senyai]

likely, will check tomorrow

[senyai]

Thank you, @hpkfft! Using nb::cast, I got the two different types I expected: dtype=int32 and dtype=float32. I initially misunderstood the documentation, but now I see where it explains that this is possible.

[hpkfft]

If you write a C++ template instead, see https://nanobind.readthedocs.io/en/latest/functions.html#function-templates