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

Author: 9prady9

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.

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

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():

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):
     """

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))