Extend dpctl_ext.tensor with the remaining functions (#2806)

This PR extends `dpctl_ext.tensor` API with the remaining statistical
and testing functions adding `std(), var(), mean(), allclose()`

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -179,6 +179,8 @@
     unique_values,
 )
 from ._sorting import argsort, sort, top_k
+from ._statistical_functions import mean, std, var
+from ._testing import allclose
 from ._type_utils import can_cast, finfo, iinfo, isdtype, result_type
 from ._utility_functions import all, any, diff
 
@@ -188,6 +190,7 @@
     "acosh",
     "add",
     "all",
+    "allclose",
     "angle",
     "any",
     "arange",
@@ -267,6 +270,7 @@
     "log10",
     "max",
     "maximum",
+    "mean",
     "meshgrid",
     "min",
     "minimum",
@@ -308,6 +312,7 @@
     "square",
     "squeeze",
     "stack",
+    "std",
     "subtract",
     "sum",
     "swapaxes",
@@ -327,6 +332,7 @@
     "unique_inverse",
     "unique_values",
     "unstack",
+    "var",
     "vecdot",
     "where",
     "zeros",

dpctl_ext/tensor/_clip.py

@@ -28,12 +28,12 @@
 
 import dpctl
 import dpctl.tensor as dpt
-import dpctl.tensor._tensor_elementwise_impl as tei
 from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
 
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
 import dpctl_ext.tensor as dpt_ext
+import dpctl_ext.tensor._tensor_elementwise_impl as tei
 import dpctl_ext.tensor._tensor_impl as ti
 
 from ._copy_utils import (

dpctl_ext/tensor/_ctors.py

@@ -361,7 +361,7 @@ def _copy_through_host_walker(seq_o, usm_res):
             )
             is None
         ):
-            usm_res[...] = dpt.asnumpy(seq_o).copy()
+            usm_res[...] = dpt_ext.asnumpy(seq_o).copy()
             return
         else:
             usm_res[...] = seq_o

dpctl_ext/tensor/_reduction.py

@@ -506,7 +506,7 @@ def count_nonzero(x, /, *, axis=None, keepdims=False, out=None):
             type.
     """
     if x.dtype != dpt.bool:
-        x = dpt.astype(x, dpt.bool, copy=False)
+        x = dpt_ext.astype(x, dpt.bool, copy=False)
     return sum(
         x,
         axis=axis,

dpctl_ext/tensor/_set_functions.py

@@ -30,11 +30,11 @@
 
 import dpctl.tensor as dpt
 import dpctl.utils as du
-from dpctl.tensor._tensor_elementwise_impl import _not_equal, _subtract
 
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
 import dpctl_ext.tensor as dpt_ext
+from dpctl_ext.tensor._tensor_elementwise_impl import _not_equal, _subtract
 
 from ._copy_utils import _empty_like_orderK
 from ._scalar_utils import (