Writing Your Own UFunc and Generalized UFunc For a Custom Data Type

The following are several examples for creating your own ufuncs and generalized

ufuncs for custom data types. All the examples use the custom dtype 'Rational'

located in the numpy-dtypes repository on github.com

Extending Existing UFunc for Custom DType

The first example shows how to extend the existing 'add' ufunc for the Rational

dtype. The add ufunc is extended for Rational dtypes using Rational's

'rational_add' function which takes two rational numbers and returns a rational

object representing the sum of those two rational numbers:

static NPY_INLINE rational

rational_add(rational x, rational y) {

// d(y) retrieves denominator for y

return make_rational_fast((int64_t)x.n*d(y) + (int64_t)d(x)*y.n, (int64_t)d(x)*d(y));

}

1. A 1-d loop function is created which loops over each pair of elements from two

1-d arrays of rationals and calls rational_add for each pair:

void rational_ufunc_add(char** args, npy_intp* dimensions, npy_intp* steps, void* data) {

npy_intp is0 = steps[0], is1 = steps[1], os = steps[2], n = *dimensions;

char *i0 = args[0], *i1 = args[1], *o = args[2];

int k;

for (k = 0; k < n; k++) {

rational x = *(rational*)i0;

rational y = *(rational*)i1;

*(rational*)o = rational_add(x,y);

i0 += is0; i1 += is1; o += os;

}

}

The loop function must have the exact signature as above. The function parameters are:

char **args - array of pointers to the actual data for the input and output arrays.

In this example there are three pointers: two pointers pointing to blocks of

memory for the two input arrays, and one pointer pointing to the block of

memory for the output array. The result of rational_add should be stored in

the output array.

dimensions - a pointer to the size of the dimension over which this function is looping

steps - a pointer to the number of bytes to jump to get to the next element in this

dimension for each of the input and output arguments

data - arbitrary data (extra arguments, function names, etc.) that can be stored with

the ufunc and will be passed in when it is called

The Rational dtype has an example of a C macro which can be used to generate create the

above function for different rational ufuncs:

#define BINARY_UFUNC(name,intype0,intype1,outtype,exp) \

void name(char** args, npy_intp* dimensions, npy_intp* steps, void* data) { \

npy_intp is0 = steps[0], is1 = steps[1], os = steps[2], n = *dimensions; \

char *i0 = args[0], *i1 = args[1], *o = args[2]; \

int k; \

for (k = 0; k < n; k++) { \

intype0 x = *(intype0*)i0; \

intype1 y = *(intype1*)i1; \

*(outtype*)o = exp; \

i0 += is0; i1 += is1; o += os; \

} \

}

#define RATIONAL_BINARY_UFUNC(name, type, exp) BINARY_UFUNC(rational_ufunc_##name, rational, rational, type, exp)

which can be used like so:

RATIONAL_BINARY_UFUNC(add, rational, rational_add(x,y))

with the following arguments:

- name suffix of 1-d loop function (the generated loop function will have the name

rational_ufunc_<name>'

- output type

- expression to calculate the output value for each pair of input elements. In this

example the expression is a call to the function rational_add.

2. In the 'initrational' function used to initialize the Rational dtype with numpy, a

PyUFuncObject is obtained for the existing 'add' ufunc in the numpy module:

PyUFuncObject* ufunc = (PyUFuncObject*)PyObject_GetAttrString(numpy,"add");

3. The 1-d loop function is registered using the PyUFuncObject obtained in step 2:

int types[] = {npy_rational,npy_rational,npy_rational};

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_add,types,0) < 0) {

return;

}

The function parameters are:

- pointer to PyUFuncObject obtained in step 2

- custom rational dtype id (obtained when dtype is registered with call to PyArray_RegisterDataType)

- 1-d loop function

- array of input and output type ids (in this case two input rational types and one

output rational type)

- pointer to arbitrary data that will be passed to 1-d loop function

4. Steps 2-3 can also be accomplished by using a c MACRO similar to the one

provided with Rational:

#define REGISTER_UFUNC(name,...) { \

PyUFuncObject* ufunc = (PyUFuncObject*)PyObject_GetAttrString(numpy,#name); \

if (!ufunc) { \

return; \

} \

int _types[] = __VA_ARGS__; \

if (sizeof(_types)/sizeof(int)!=ufunc->nargs) { \

PyErr_Format(PyExc_AssertionError,"ufunc %s takes %d arguments, our loop takes %ld",#name,ufunc->nargs,sizeof(_types)/sizeof(int)); \

return; \

} \

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_##name,_types,0)<0) { \

return; \

} \

}

#define REGISTER_UFUNC_BINARY_RATIONAL(name) REGISTER_UFUNC(name,{npy_rational,npy_rational,npy_rational})

REGISTER_UFUNC_BINARY_RATIONAL(add)

Creating New UFunc for Custom DType

The next example shows how to create a new ufunc for the Rational dtype. The ufunc

example creates a ufunc called 'numerator' which generates an array of numerator

values based rational numbers from the input array.

1. A 1-d loop function is created as before which takes the numerator value from

each element of the input array and stores it in the output array:

void rational_ufunc_numerator(char** args, npy_intp* dimensions, npy_intp* steps, void* data) {

npy_intp is = steps[0], os = steps[1], n = *dimensions;

char *i = args[0], *o = args[1];

int k;

for (k = 0; k < n; k++) {

rational x = *(rational*)i;

*(int64_t*)o = x.n;

i += is; o += os;

}

}

You can also use the c MACRO provided in Rational for generating the above function:

#define UNARY_UFUNC(name,type,exp) \

void rational_ufunc_##name(char** args, npy_intp* dimensions, npy_intp* steps, void* data) { \

npy_intp is = steps[0], os = steps[1], n = *dimensions; \

char *i = args[0], *o = args[1]; \

int k; \

for (k = 0; k < n; k++) { \

rational x = *(rational*)i; \

*(type*)o = exp; \

i += is; o += os; \

} \

}

UNARY_UFUNC(numerator,int64_t,x.n)

2. In the 'initrational' function used to initialize the Rational dtype with numpy,

a new PyUFuncObject is created for the new 'numerator' ufunc using the

PyUFunc_FromFuncAndData function:

PyObject* ufunc = PyUFunc_FromFuncAndData(0,0,0,0,1,1,PyUFunc_None,(char*)"numerator",(char*)"rational number numerator",0);

In this use case, the first four parameters should be set to zero since we're

creating a new ufunc without support for existing dtypes. The rest of the

parameters:

- number of inputs to function that the loop function calls for each pair of elements

- number of outputs of loop function

- name of the ufunc

- documentation string describing the ufunc

- unused; present for backwards compatibility

3. The 1-d loop function is registered using the loop function and the

PyUFuncObject created in step 2:

int _types[] = {npy_rational,NPY_INT64};

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_numerator,_types,0)<0) {

return;

}

4. Finally, a function called 'numerator' is added to the rational module which

will call the numerator ufunc:

PyModule_AddObject(m,"numerator",(PyObject*)ufunc);

5. Steps 2-4 can also be accomplished by using a c MACRO similar to the one

provided with Rational:

#define NEW_UNARY_UFUNC(name,type,doc) { \

PyObject* ufunc = PyUFunc_FromFuncAndData(0,0,0,0,1,1,PyUFunc_None,(char*)#name,(char*)doc,0); \

if (!ufunc) { \

return; \

} \

int types[2] = {npy_rational,type}; \

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_##name,types,0)<0) { \

return; \

} \

PyModule_AddObject(m,#name,(PyObject*)ufunc); \

}

NEW_UNARY_UFUNC(numerator,NPY_INT64,"rational number numerator");

Creating New Generalized UFunc for Custom DType

The next example shows how to create a new generalized ufunc for the Rational dtype.

The gufunc example creates a gufunc 'matrix_multiply' which loops over a pair of

vectors or matrices and performs a matrix multiply on each pair of matrix elements.

1. A loop function is created to loop through the outer or loop dimensions, performing a

matrix multiply operation on the core dimensions for each loop:

static void

rational_gufunc_matrix_multiply(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))

{

// outer dimensions counter

npy_intp N_;

// length of flattened outer dimensions

npy_intp dN = dimensions[0];

// striding over flattened outer dimensions for input and output arrays

npy_intp s0 = steps[0];

npy_intp s1 = steps[1];

npy_intp s2 = steps[2];

// loop through outer dimensions, performing matrix multiply on core dimensions for each loop

for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2) {

rational_matrix_multiply(args, dimensions+1, steps+3);

}

}

If the input matrices have more than one outer dimension, the outer dimensions are flattened from the

perspective of the loop function.

The matrix multiply function:

static NPY_INLINE void

rational_matrix_multiply(char **args, npy_intp *dimensions, npy_intp *steps)

{

// pointers to data for input and output arrays

char *ip1 = args[0];

char *ip2 = args[1];

char *op = args[2];

// lengths of core dimensions

npy_intp dm = dimensions[0];

npy_intp dn = dimensions[1];

npy_intp dp = dimensions[2];

// striding over core dimensions

npy_intp is1_m = steps[0];

npy_intp is1_n = steps[1];

npy_intp is2_n = steps[2];

npy_intp is2_p = steps[3];

npy_intp os_m = steps[4];

npy_intp os_p = steps[5];

// core dimensions counters

npy_intp m, n, p;

// calculate dot product for each row/column vector pair

for (m = 0; m < dm; m++) {

for (p = 0; p < dp; p++) {

npyrational_dot(ip1, is1_n, ip2, is2_n, op, dn, NULL);

ip2 += is2_p;

op += os_p;

}

ip2 -= is2_p * p;

op -= os_p * p;

ip1 += is1_m;

op += os_m;

}

}

2. In the 'initrational' function used to initialize the Rational dtype with numpy,

a new PyUFuncObject is created for the new 'matrix_multiply' generalized ufunc using the

PyUFunc_FromFuncAndDataAndSignature function:

PyObject* gufunc = PyUFunc_FromFuncAndDataAndSignature(0,0,0,0,2,1,PyUFunc_None,(char*)"matrix_multiply",(char*)"return result of multiplying two matrices of rationals",0,"(m,n),(n,p)->(m,p)");

This is identical to the PyUFunc_FromFuncAndData function used to create a ufunc object in the examples above,

with the addition of a ufunc signature describing the core dimensions of the input and output arrays. In this example,

the generalized ufunc operates on pairs of matrices with dimensions (m,n) and (n,p), producing an

output matrix of dimensions (m,p).

3. The loop function is registered using the loop function and the PyUFuncObject created in step 2:

int types2[3] = {npy_rational,npy_rational,npy_rational};

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)gufunc,npy_rational,rational_gufunc_matrix_multiply,types2,0) < 0) {

return;

}

4. Finally, a function called 'matrix_multiply' is added to the rational module which

will call the numerator ufunc:

PyModule_AddObject(m,"matrix_multiply",(PyObject*)gufunc);