From ae6899186ab1e4b874f50a819d44f635750ee725 Mon Sep 17 00:00:00 2001 From: pradeep Date: Wed, 3 Mar 2021 17:02:20 +0530 Subject: [PATCH] Refactor cov, var, stdev APIs to reflect 3.8 release Also adds set_cublas_mode utility function to explicitly enable tensor ops for blas functions --- arrayfire/algorithm.py | 28 ++++++++++++++++++++ arrayfire/cuda.py | 7 +++++ arrayfire/library.py | 7 +++++ arrayfire/statistics.py | 53 ++++++++++++++++++-------------------- tests/simple/statistics.py | 8 +++--- 5 files changed, 71 insertions(+), 32 deletions(-) diff --git a/arrayfire/algorithm.py b/arrayfire/algorithm.py index d5adbcce5..6f06cee14 100644 --- a/arrayfire/algorithm.py +++ b/arrayfire/algorithm.py @@ -250,6 +250,34 @@ def max(a, dim=None): else: return _reduce_all(a, backend.get().af_max_all) +def maxRagged(vals, lens, dim): + """ + Find the maximum value of a subset of elements along a specified dimension. + + The size of the subset of elements along the given dimension are decided based on the lengths + provided in the `lens` array. + + Parameters + ---------- + vals : af.Array + Multi dimensional arrayfire array. + lens : af.Array + Multi dimensional arrayfire array containing number of elements to reduce along given `dim` + dim: optional: int. default: None + Dimension along which the maximum value is required. + + Returns + ------- + (values, indices): A tuple of af.Array(s) + `values` af.Array will have the maximum values along given dimension for + subsets determined by lengths provided in `lens` + `idx` contains the locations of the maximum values as per the lengths provided in `lens` + """ + out_vals = Array() + out_idx = Array() + safe_call(backend().get().af_max_ragged(c_pointer(out_vals.arr), c_pointer(out_idx.arr), c_pointer(vals.arr), c_pointer(lens.arr), c_int_t(dim))) + return out_vals, out_idx + def maxByKey(keys, vals, dim=-1): """ Calculate the max of elements along a specified dimension according to a key. diff --git a/arrayfire/cuda.py b/arrayfire/cuda.py index ea24c0e30..e5ef819be 100644 --- a/arrayfire/cuda.py +++ b/arrayfire/cuda.py @@ -85,3 +85,10 @@ def set_native_id(idx): safe_call(backend.get().afcu_set_native_id(idx)) return + +def set_cublas_mode(mode=CUBLAS_MATH_MODE.DEFAULT): + """ + Set's cuBLAS math mode for CUDA backend. In other backends, this has no effect. + """ + safe_call(backend().get().afcu_cublasSetMathMode(mode.value)) + return diff --git a/arrayfire/library.py b/arrayfire/library.py index 27ca4f4d9..1b3c8b3ea 100644 --- a/arrayfire/library.py +++ b/arrayfire/library.py @@ -490,6 +490,13 @@ class VARIANCE(_Enum): SAMPLE = _Enum_Type(1) POPULATION = _Enum_Type(2) +class CUBLAS_MATH_MODE(_Enum): + """ + Enable Tensor Core usage if available on CUDA backend GPUs + """ + DEFAULT = _Enum_Type(0) + TENSOR_OP = _Enum_Type(1) + _VER_MAJOR_PLACEHOLDER = "__VER_MAJOR__" def _setup(): diff --git a/arrayfire/statistics.py b/arrayfire/statistics.py index f47f3a48d..158da18de 100644 --- a/arrayfire/statistics.py +++ b/arrayfire/statistics.py @@ -59,7 +59,7 @@ def mean(a, weights=None, dim=None): return real if imag == 0 else real + imag * 1j -def var(a, isbiased=False, weights=None, dim=None): +def var(a, bias=VARIANCE.DEFAULT, weights=None, dim=None): """ Calculate variance along a given dimension. @@ -68,9 +68,9 @@ def var(a, isbiased=False, weights=None, dim=None): a: af.Array The input array. - isbiased: optional: Boolean. default: False. - Boolean denoting population variance (false) or sample - variance (true). + bias: optional: af.VARIANCE. default: DEFAULT. + population variance(VARIANCE.POPULATION) or sample variance(VARIANCE.SAMPLE). + This is ignored if weights are provided. weights: optional: af.Array. default: None. Array to calculate for the weighted mean. Must match size of @@ -89,7 +89,7 @@ def var(a, isbiased=False, weights=None, dim=None): out = Array() if weights is None: - safe_call(backend.get().af_var(c_pointer(out.arr), a.arr, isbiased, c_int_t(dim))) + safe_call(backend.get().af_var_v2(c_pointer(out.arr), a.arr, bias.value, c_int_t(dim))) else: safe_call(backend.get().af_var_weighted(c_pointer(out.arr), a.arr, weights.arr, c_int_t(dim))) @@ -99,7 +99,7 @@ def var(a, isbiased=False, weights=None, dim=None): imag = c_double_t(0) if weights is None: - safe_call(backend.get().af_var_all(c_pointer(real), c_pointer(imag), a.arr, isbiased)) + safe_call(backend.get().af_var_all_v2(c_pointer(real), c_pointer(imag), a.arr, bias.value)) else: safe_call(backend.get().af_var_all_weighted(c_pointer(real), c_pointer(imag), a.arr, weights.arr)) @@ -150,7 +150,7 @@ def meanvar(a, weights=None, bias=VARIANCE.DEFAULT, dim=-1): return mean_out, var_out -def stdev(a, dim=None): +def stdev(a, bias=VARIANCE.DEFAULT, dim=None): """ Calculate standard deviation along a given dimension. @@ -159,6 +159,10 @@ def stdev(a, dim=None): a: af.Array The input array. + bias: optional: af.VARIANCE. default: DEFAULT. + population variance(VARIANCE.POPULATION) or sample variance(VARIANCE.SAMPLE). + This is ignored if weights are provided. + dim: optional: int. default: None. The dimension for which to obtain the standard deviation from input data. @@ -171,48 +175,41 @@ def stdev(a, dim=None): """ if dim is not None: out = Array() - safe_call(backend.get().af_stdev(c_pointer(out.arr), a.arr, c_int_t(dim))) + safe_call(backend.get().af_stdev_v2(c_pointer(out.arr), a.arr, bias.value, + c_int_t(dim))) return out else: real = c_double_t(0) imag = c_double_t(0) - safe_call(backend.get().af_stdev_all(c_pointer(real), c_pointer(imag), a.arr)) + safe_call(backend.get().af_stdev_all_v2(c_pointer(real), c_pointer(imag), a.arr, + bias.value)) real = real.value imag = imag.value return real if imag == 0 else real + imag * 1j -def cov(a, isbiased=False, dim=None): +def cov(a, b, bias=VARIANCE.DEFAULT): """ Calculate covariance along a given dimension. Parameters ---------- a: af.Array - The input array. + Input array. - isbiased: optional: Boolean. default: False. - Boolean denoting whether biased estimate should be taken. + b: af.Array + Input array. - dim: optional: int. default: None. - The dimension for which to obtain the covariance from input data. + bias: optional: af.VARIANCE. default: DEFAULT. + population variance(VARIANCE.POPULATION) or sample variance(VARIANCE.SAMPLE). Returns ------- output: af.Array - Array containing the covariance of the input array along a - given dimension. + Array containing the covariance of the input array along a given dimension. """ - if dim is not None: - out = Array() - safe_call(backend.get().af_cov(c_pointer(out.arr), a.arr, isbiased, c_int_t(dim))) - return out - else: - real = c_double_t(0) - imag = c_double_t(0) - safe_call(backend.get().af_cov_all(c_pointer(real), c_pointer(imag), a.arr, isbiased)) - real = real.value - imag = imag.value - return real if imag == 0 else real + imag * 1j + out = Array() + safe_call(backend.get().af_cov_v2(c_pointer(out.arr), a.arr, b.arr, bias.value)) + return out def median(a, dim=None): """ diff --git a/tests/simple/statistics.py b/tests/simple/statistics.py index 174af0a5b..39fe6703f 100644 --- a/tests/simple/statistics.py +++ b/tests/simple/statistics.py @@ -28,10 +28,10 @@ def simple_statistics(verbose=False): print_func(af.mean(a, weights=w)) display_func(af.var(a, dim=0)) - display_func(af.var(a, isbiased=True, dim=0)) + display_func(af.var(a, bias=af.VARIANCE.SAMPLE, dim=0)) display_func(af.var(a, weights=w, dim=0)) print_func(af.var(a)) - print_func(af.var(a, isbiased=True)) + print_func(af.var(a, bias=af.VARIANCE.SAMPLE)) print_func(af.var(a, weights=w)) mean, var = af.meanvar(a, dim=0) @@ -45,9 +45,9 @@ def simple_statistics(verbose=False): print_func(af.stdev(a)) display_func(af.var(a, dim=0)) - display_func(af.var(a, isbiased=True, dim=0)) + display_func(af.var(a, bias=af.VARIANCE.SAMPLE, dim=0)) print_func(af.var(a)) - print_func(af.var(a, isbiased=True)) + print_func(af.var(a, bias=af.VARIANCE.SAMPLE)) display_func(af.median(a, dim=0)) print_func(af.median(w))