allow inf or NaN cval in ndimage.filter functions as well
@@ -142,6 +142,9 @@ def _call_kernel(kernel, input, weights, output, structure=None,
return output
+math_constants_preamble = "#include <math_constants.h>\n"
+
+
_CAST_FUNCTION = """
// Implements a casting function to make it compatible with scipy
// Use like cast<to_type>(value)
@@ -252,6 +255,13 @@ def _generate_nd_kernel(name, pre, found, post, mode, w_shape, int_type,
if mode == 'constant':
cond = ' || '.join(['(ix_{} < 0)'.format(j) for j in range(ndim)])
+ if cval is numpy.nan:
+ cval = 'CUDART_NAN'
+ elif cval == numpy.inf:
+ cval = 'CUDART_INF'
+ elif cval == -numpy.inf:
+ cval = '-CUDART_INF'
+
if binary_morphology:
found = found.format(cond=cond, value=value)
else:
@@ -290,7 +300,7 @@ def _generate_nd_kernel(name, pre, found, post, mode, w_shape, int_type,
name += '_with_structure'
if has_mask:
name += '_with_mask'
- preamble = _CAST_FUNCTION + preamble
+ preamble = math_constants_preamble + _CAST_FUNCTION + preamble
return cupy.ElementwiseKernel(in_params, out_params, operation, name,
reduce_dims=False, preamble=preamble,
options=('--std=c++11',) + options)
