| 838 | |
| 839 | template <typename UFunc> |
| 840 | bool RegisterUFunc(PyObject* numpy, const char* name) { |
| 841 | std::vector<int> types = UFunc::Types(); |
| 842 | PyUFuncGenericFunction fn = UFunc::Call; |
| 843 | Safe_PyObjectPtr ufunc_obj = make_safe(PyObject_GetAttrString(numpy, name)); |
| 844 | if (!ufunc_obj) { |
| 845 | return false; |
| 846 | } |
| 847 | PyUFuncObject* ufunc = reinterpret_cast<PyUFuncObject*>(ufunc_obj.get()); |
| 848 | if (static_cast<int>(types.size()) != ufunc->nargs) { |
| 849 | PyErr_Format(PyExc_AssertionError, |
| 850 | "ufunc %s takes %d arguments, loop takes %lu", name, |
| 851 | ufunc->nargs, types.size()); |
| 852 | return false; |
| 853 | } |
| 854 | if (PyUFunc_RegisterLoopForType(ufunc, npy_bfloat16, fn, |
| 855 | const_cast<int*>(types.data()), |
| 856 | nullptr) < 0) { |
| 857 | return false; |
| 858 | } |
| 859 | return true; |
| 860 | } |
| 861 | |
| 862 | namespace ufuncs { |
| 863 | |