How to make cython function accept float or double array input?

Question

Suppose I have the following (MCVE...) cython function

cimport cython

from scipy.linalg.cython_blas cimport dnrm2


cpdef double func(int n, double[:] x):
   cdef int inc = 1
   return dnrm2(&n, &x[0], &inc)

Then, I cannot call it on a np.float32 array x.

How could I make func accept a double[:] or a float[:], and call dnrm2 or snrm2 alternatively? The only solution I have currently is to have two functions, which creates a huge quantity of duplicated code.

Answer 1

You could use a fused type. Please note that the below doesn't compile on my system because ddot and sdot apparently require 5 parameters:

# cython: infer_types=True
cimport cython

from scipy.linalg.cython_blas cimport ddot, sdot

ctypedef fused anyfloat:
   double
   float

cpdef anyfloat func(int n, anyfloat[:] x):
   cdef int inc = 1
   if anyfloat is double:
      return ddot(&n, &x[0], &inc)
   else:
      return sdot(&n, &x[0], &inc)