Initializes the module.
| 585 | |
| 586 | // Initializes the module. |
| 587 | bool Initialize() { |
| 588 | // It's critical to ImportNumpy and import umath |
| 589 | // to avoid crash in open source build. |
| 590 | ImportNumpy(); |
| 591 | import_umath1(false); |
| 592 | |
| 593 | Safe_PyObjectPtr numpy_str = make_safe(MakePyString("numpy")); |
| 594 | if (!numpy_str) { |
| 595 | return false; |
| 596 | } |
| 597 | Safe_PyObjectPtr numpy = make_safe(PyImport_Import(numpy_str.get())); |
| 598 | if (!numpy) { |
| 599 | return false; |
| 600 | } |
| 601 | |
| 602 | // We hit a mysterious crash if we haven't initialized numpy before this: |
| 603 | PyBfloat16_Type.tp_base = &PyGenericArrType_Type; |
| 604 | |
| 605 | if (PyType_Ready(&PyBfloat16_Type) < 0) { |
| 606 | return false; |
| 607 | } |
| 608 | |
| 609 | // Initializes the NumPy descriptor. |
| 610 | PyArray_InitArrFuncs(&NPyBfloat16_ArrFuncs); |
| 611 | NPyBfloat16_ArrFuncs.getitem = NPyBfloat16_GetItem; |
| 612 | NPyBfloat16_ArrFuncs.setitem = NPyBfloat16_SetItem; |
| 613 | NPyBfloat16_ArrFuncs.compare = NPyBfloat16_Compare; |
| 614 | NPyBfloat16_ArrFuncs.copyswapn = NPyBfloat16_CopySwapN; |
| 615 | NPyBfloat16_ArrFuncs.copyswap = NPyBfloat16_CopySwap; |
| 616 | NPyBfloat16_ArrFuncs.nonzero = NPyBfloat16_NonZero; |
| 617 | NPyBfloat16_ArrFuncs.fill = NPyBfloat16_Fill; |
| 618 | |
| 619 | Py_TYPE(&NPyBfloat16_Descr) = &PyArrayDescr_Type; |
| 620 | npy_bfloat16_ = PyArray_RegisterDataType(&NPyBfloat16_Descr); |
| 621 | if (npy_bfloat16_ < 0) return false; |
| 622 | |
| 623 | // Support dtype(bfloat16) |
| 624 | if (PyDict_SetItemString(PyBfloat16_Type.tp_dict, "dtype", |
| 625 | reinterpret_cast<PyObject*>(&NPyBfloat16_Descr)) < |
| 626 | 0) { |
| 627 | return false; |
| 628 | } |
| 629 | |
| 630 | // Register casts |
| 631 | |
| 632 | // We lie shamelessly and say that a cast from half to bfloat16 is safe. |
| 633 | // Numpy frequently uses the smallest legal representation type for small |
| 634 | // float constants (e.g., 1.0), which is often float16. Things break if these |
| 635 | // cannot be converted transparently to bfloat16. |
| 636 | if (!RegisterBfloat16Cast<Eigen::half>(NPY_HALF, /*cast_is_safe=*/true)) { |
| 637 | return false; |
| 638 | } |
| 639 | |
| 640 | if (!RegisterBfloat16Cast<float>(NPY_FLOAT, /*cast_is_safe=*/true)) { |
| 641 | return false; |
| 642 | } |
| 643 | if (!RegisterBfloat16Cast<double>(NPY_DOUBLE, /*cast_is_safe=*/true)) { |
| 644 | return false; |