C-API register a ufunc loop against builtin ufunc

Question

I am looking to build a numpy extension module which declares a new structured dtype and provides the necessary inner ufunc loops to allow built-in math operations to be performed on it.

Following the guide on creating your own ufunc, I have been able to:

• Create a new PyUFuncGenericFunction to mock a built-in ufunc:

  PyUFuncGenericFunction *add_uncertain = PyUFunc_FromFuncAndData(
      NULL,
      NULL,
      NULL,
      0,
      2,
      1,
      PyUFunc_None,
      "add_uncertain",
      "adds arrays of uncertain values",
      0);
  

• Declare the data-type, as:

  PyObject *uncertain_double_dtype_dict = Py_BuildValue(
      "[(s, s), (s, s)]",
      "nominal", "f8", "uncertainty", "f8");
  

• Write an implementation of this for a specific dtype:

  static PyObject *
  add_uncertain_double(char **args, npy_intp *dimensions, npy_intp *steps, void *data)
  {
      char *in_arr_a = args[0];
      char *in_arr_b = args[1];
      char *out_arr = args[2];
  
      npy_intp in_a_step = steps[0];
      npy_intp in_b_step = steps[1];
      npy_intp out_step = steps[2];
  
      for (npy_intp i = 0; i < dimensions[0]; i++)
      {
          double *in_a = (double *)in_arr_a;
          double *in_b = (double *)in_arr_b;
          double *out = (double *)out_arr;
  
          out[0] = in_a[0] + in_b[0];
          out[1] = hypot(in_a[1], in_b[1]);
  
          in_arr_a += in_a_step;
          out_arr += out_step;
      }
  };
  

• Register this function against the mock ufunc:

  PyUFunc_RegisterLoopForDescr(
      add_uncertain,
      uncertain_double_dtype,
      &add_uncertain_double,
      add_uncertain_double_dtypes,
      NULL);
  

When exported as a module this works as intended, but requires the user to call the function via the package exported ufunc (i.e. my_package.add_uncertain); I would rather this implementation was available via numpy.add.

The language used in the registering a ufunc loop section of the numpy C-API docs seems to suggest that I should be able to register the ufunc loop against built-in ufuncs. To do so I believe I should pass the built-in PyUFuncGenericFunction to PyUFunc_RegisterLoopForDescr.

I would be very grateful to know if I am on the right track with this, and if so where I should look for the built-in PyUFuncGenericFunction

Answer 1

The solution, as contained within the numpy rational type test, was to import numpy with PyImport_Import then get the add ufunc using PyObject_GetAttrString, at which point the new inner loop could be registered:

  numpy_str = PyUnicode_FromString("numpy");
  if (!numpy_str) 
      return NULL;
  numpy = PyImport_Import(numpy_str);
  Py_DECREF(numpy_str);
  if (!numpy)
      return NULL;

  PyUFuncObject *add_ufunc = PyObject_GetAttrString(numpy, "add");
  if (!add_ufunc)
      return NULL;

  PyUFunc_RegisterLoopForDescr(
      add_ufunc,
      uncertain_double_dtype,
      &add_uncertain_double,
      add_uncertain_double_dtypes,
      NULL);