Using ml_dtypes.bfloat16 with nb::ndarray

[awni]

Using bfloat16 from ml_dtypes as the type of a np.ndarray is not cast to a nd::array.

Is there a way to use the non-standard types to get this to work?

For example, given this extension:

#include <nanobind/nanobind.h>                                                 
#include <nanobind/ndarray.h>                                                  
                                                                               
namespace nb = nanobind;                                                       
                                                                               
void check(nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>) {               
  std::cout << "Got ND Array? " << std::endl;                                  
}                                                                              
                                                                               
NB_MODULE(ext, m) {                                                            
    m.def("check", &check);                                                    
}       

The following code throws an error on the bfloat array (which is expected).

import ext
from ml_dtypes import bfloat16
import numpy as np

x = np.array(1., dtype=np.float16)
ext.check(x)

x = np.array(1., dtype=bfloat16)
# Error:
# TypeError: check(): incompatible function arguments. The following argument types are supported:
#    1. check(arg: typing.Union[float, ndarray[writable=False, order='C', device='cpu']], /) -> None
#
# Invoked with types: ndarray
ext.check(x)

[wjakob]

What would be the C++ type associated with bfloat? If you have that, then it is likely possible to use the documented approach.

[awni]

That's what I'm looking for now. We have a bfloat16 in MLX already that is registered as a custom type with nb::ndarray: https://github.com/ml-explore/mlx/blob/main/python/src/convert.cpp#L19-L26

I am trying to figure out how we can map the ml_dtype.bfloat16 to the MLX bfloat16 type.

[awni]

That's the ml_dtype bfloat, it looks like it's using Eigen: https://github.com/jax-ml/ml_dtypes/blob/main/ml_dtypes/_src/dtypes.cc#L40-L60

[awni]

I'm not sure I follow how the documented approach will work here. It looks like nanobind doesn't cast the incoming np.ndarray to a nb::ndarray (but it casts fine for fp16 which is also non-standard). If I register the Eigen::bfloat16 with Nanobind would that cause the cast to succeed?

namespace nanobind {
    template <> struct ndarray_traits<Eigen::bfloat16> {
        static constexpr bool is_complex = false;
        static constexpr bool is_float   = true;
        static constexpr bool is_bool    = false;
        static constexpr bool is_int     = false;
        static constexpr bool is_signed  = true;
    };
}