| 287 | } |
| 288 | |
| 289 | absl::optional<CastToArrayResult> CastToArray(py::handle h) { |
| 290 | py::array array = py::array::ensure( |
| 291 | h, py::array::c_style | py::detail::npy_api::NPY_ARRAY_ALIGNED_); |
| 292 | if (!array) { |
| 293 | return absl::nullopt; |
| 294 | } |
| 295 | |
| 296 | auto type_or_status = DtypeToPrimitiveType(array.dtype()); |
| 297 | if (!type_or_status.ok()) { |
| 298 | throw std::runtime_error(type_or_status.status().ToString()); |
| 299 | } |
| 300 | PrimitiveType type = type_or_status.ValueOrDie(); |
| 301 | |
| 302 | if (type == BF16) { |
| 303 | // The NumPy array protocol has no way to describe our custom bfloat16 |
| 304 | // type, so we cast to an array of uint16 instead. We are going to pass |
| 305 | // a raw buffer pointer to BorrowingLiteral anyway, so it doesn't |
| 306 | // really matter what type we use here, so long as it has the correct size |
| 307 | // and alignment. |
| 308 | array = py::reinterpret_steal<py::array>( |
| 309 | PyArray_View(reinterpret_cast<PyArrayObject*>(array.ptr()), |
| 310 | reinterpret_cast<PyArray_Descr*>( |
| 311 | py::dtype::of<uint16>().release().ptr()), |
| 312 | static_cast<PyTypeObject*>(nullptr))); |
| 313 | } |
| 314 | |
| 315 | py::buffer_info buffer_info = array.request(); |
| 316 | |
| 317 | absl::InlinedVector<int64, 4> dims(array.ndim()); |
| 318 | for (int i = 0; i < array.ndim(); ++i) { |
| 319 | dims[i] = array.shape(i); |
| 320 | } |
| 321 | Shape shape = ShapeUtil::MakeShape(type, dims); |
| 322 | if (buffer_info.size * buffer_info.itemsize != ShapeUtil::ByteSizeOf(shape)) { |
| 323 | throw std::runtime_error(absl::StrCat( |
| 324 | "Size mismatch for buffer: ", buffer_info.size * buffer_info.itemsize, |
| 325 | " vs. ", ShapeUtil::ByteSizeOf(shape))); |
| 326 | } |
| 327 | return CastToArrayResult{array, static_cast<const char*>(buffer_info.ptr), |
| 328 | shape}; |
| 329 | } |
| 330 | |
| 331 | } // namespace xla |