Initializes the module.
| 1288 | |
| 1289 | // Initializes the module. |
| 1290 | bool Initialize() { |
| 1291 | import_array1(false); |
| 1292 | import_umath1(false); |
| 1293 | |
| 1294 | Safe_PyObjectPtr numpy_str = make_safe(MakePyString("numpy")); |
| 1295 | if (!numpy_str) { |
| 1296 | return false; |
| 1297 | } |
| 1298 | Safe_PyObjectPtr numpy = make_safe(PyImport_Import(numpy_str.get())); |
| 1299 | if (!numpy) { |
| 1300 | return false; |
| 1301 | } |
| 1302 | |
| 1303 | PyBfloat16_Type.tp_base = &PyGenericArrType_Type; |
| 1304 | |
| 1305 | if (PyType_Ready(&PyBfloat16_Type) < 0) { |
| 1306 | return false; |
| 1307 | } |
| 1308 | |
| 1309 | // Initializes the NumPy descriptor. |
| 1310 | PyArray_InitArrFuncs(&NPyBfloat16_ArrFuncs); |
| 1311 | NPyBfloat16_ArrFuncs.getitem = NPyBfloat16_GetItem; |
| 1312 | NPyBfloat16_ArrFuncs.setitem = NPyBfloat16_SetItem; |
| 1313 | NPyBfloat16_ArrFuncs.copyswapn = NPyBfloat16_CopySwapN; |
| 1314 | NPyBfloat16_ArrFuncs.copyswap = NPyBfloat16_CopySwap; |
| 1315 | NPyBfloat16_ArrFuncs.nonzero = NPyBfloat16_NonZero; |
| 1316 | NPyBfloat16_ArrFuncs.fill = NPyBfloat16_Fill; |
| 1317 | NPyBfloat16_ArrFuncs.dotfunc = NPyBfloat16_DotFunc; |
| 1318 | NPyBfloat16_ArrFuncs.compare = NPyBfloat16_CompareFunc; |
| 1319 | NPyBfloat16_ArrFuncs.argmax = NPyBfloat16_ArgMaxFunc; |
| 1320 | NPyBfloat16_ArrFuncs.argmin = NPyBfloat16_ArgMinFunc; |
| 1321 | |
| 1322 | Py_TYPE(&NPyBfloat16_Descr) = &PyArrayDescr_Type; |
| 1323 | npy_bfloat16 = PyArray_RegisterDataType(&NPyBfloat16_Descr); |
| 1324 | if (npy_bfloat16 < 0) { |
| 1325 | return false; |
| 1326 | } |
| 1327 | |
| 1328 | // Support dtype(bfloat16) |
| 1329 | if (PyDict_SetItemString(PyBfloat16_Type.tp_dict, "dtype", |
| 1330 | reinterpret_cast<PyObject*>(&NPyBfloat16_Descr)) < |
| 1331 | 0) { |
| 1332 | return false; |
| 1333 | } |
| 1334 | |
| 1335 | // Register casts |
| 1336 | if (!RegisterBfloat16Cast<Eigen::half>(NPY_HALF, /*cast_is_safe=*/false)) { |
| 1337 | return false; |
| 1338 | } |
| 1339 | if (!RegisterBfloat16Cast<float>(NPY_FLOAT, /*cast_is_safe=*/true)) { |
| 1340 | return false; |
| 1341 | } |
| 1342 | if (!RegisterBfloat16Cast<double>(NPY_DOUBLE, /*cast_is_safe=*/true)) { |
| 1343 | return false; |
| 1344 | } |
| 1345 | if (!RegisterBfloat16Cast<bool>(NPY_BOOL, /*cast_is_safe=*/false)) { |
| 1346 | return false; |
| 1347 | } |