[demo] Getting nanobind to call ndarray conversion functions from C++ after initial ingress

[ThisIsFineTM]

This is a little demo of how I was able to reach into the library to get data conversions working through nanobind after the binding interface step. In my particular use case, I needed to take a generic nb::ndarray<nb::array_api> at the binding site, but had circumstances where I needed convert the data after the fact. This effectively tells nanobind to go ask the data source library to do the conversion for us without having to fiddle with overloads at the binding site.

template <typename... Ts>
auto convert_py_tensor(nb::ndarray<Ts...> py_tensor) {

  // grab the ndarray config that has all the type/layout information
  nb::detail::ndarray_config config(typename decltype(py_tensor)::Config{});

  // convert to bfloat16
  config.dtype.code = (uint8_t)nb::dlpack::dtype_code::Bfloat;
  config.dtype.bits = 16;
  config.dtype.lanes = 1;
  
  // force row-major contiguous
  config.order = nb::c_contig::value;

  // this will call into the 3rd party library's (eg. pytorch) conversion functions
  nb::detail::ndarray_handle* handle = nb::detail::ndarray_import(
        py_tensor.cast(nb::rv_policy::automatic).ptr(),
        &config,
        true /*convert*/,
        nullptr /*(cleanup*)*/);

  // the ndarray_handle will be a nullptr if the conversion was not successful,
  // so if there was a failure in ndarray_import this ndarray would be empty.
  return nb::ndarray{handle};
}