Add cupy.cublas.iamin
@@ -104,43 +104,31 @@ def batched_gesv(a, b):
return b.transpose(0, 2, 1).reshape(b_shape)
-def _setup_result_array(handle, out, result_dtype):
- mode = cublas.getPointerMode(handle)
- if out is None or isinstance(out, cupy.ndarray):
- if out is None or out.dtype != result_dtype:
- result = cupy.empty((1,), dtype=result_dtype)
- else:
- result = out
- result_ptr = result.data.ptr
- cublas.setPointerMode(handle, cublas.CUBLAS_POINTER_MODE_DEVICE)
- elif isinstance(out, numpy.ndarray):
- if out.dtype != result_dtype:
- result = numpy.empty((1,), dtype=result_dtype)
- else:
- result = out
- result_ptr = result.ctypes.data
- cublas.setPointerMode(handle, cublas.CUBLAS_POINTER_MODE_HOST)
- else:
- raise TypeError('out must be either cupy or numpy ndarray')
- return result, result_ptr, mode
+def iamax(a, out=None):
+ return _iamaxmin(a, out, 'amax')
-def iamax(a, out=None):
+def iamin(a, out=None):
+ return _iamaxmin(a, out, 'amin')
+
+
+def _iamaxmin(a, out, name):
# Note that result index is 1-based index (not 0-based index)
if a.ndim != 1:
raise ValueError('a must be a 1D array (actual: {})'.format(a.ndim))
dtype = a.dtype.char
if dtype == 'f':
- func = cublas.isamax
+ t = 's'
elif dtype == 'd':
- func = cublas.idamax
+ t = 'd'
elif dtype == 'F':
- func = cublas.icamax
+ t = 'c'
elif dtype == 'D':
- func = cublas.izamax
+ t = 'z'
else:
raise TypeError('invalid dtype')
+ func = getattr(cublas, 'i' + t + name)
handle = device.get_cublas_handle()
result_dtype = 'i'
@@ -156,6 +144,27 @@ def iamax(a, out=None):
return out
+def _setup_result_array(handle, out, result_dtype):
+ mode = cublas.getPointerMode(handle)
+ if out is None or isinstance(out, cupy.ndarray):
+ if out is None or out.dtype != result_dtype:
+ result = cupy.empty([], dtype=result_dtype)
+ else:
+ result = out
+ result_ptr = result.data.ptr
+ cublas.setPointerMode(handle, cublas.CUBLAS_POINTER_MODE_DEVICE)
+ elif isinstance(out, numpy.ndarray):
+ if out.dtype != result_dtype:
+ result = numpy.empty([], dtype=result_dtype)
+ else:
+ result = out
+ result_ptr = result.ctypes.data
+ cublas.setPointerMode(handle, cublas.CUBLAS_POINTER_MODE_HOST)
+ else:
+ raise TypeError('out must be either cupy or numpy ndarray')
+ return result, result_ptr, mode
+
+
def asum(a, out=None):
if a.ndim != 1:
raise ValueError('a must be a 1D array (actual: {})'.format(a.ndim))
