Move _resolve_one_strong_one_weak_types and _resolve_one_strong_two_weak_types to _type_utils

Author: ndgrigorian

dpctl/tensor/_clip.py

@@ -30,124 +30,15 @@
     _validate_dtype,
 )
 from dpctl.tensor._manipulation_functions import _broadcast_shape_impl
-from dpctl.tensor._type_utils import _can_cast, _to_device_supported_dtype
+from dpctl.tensor._type_utils import _can_cast
 from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
 
 from ._type_utils import (
-    WeakComplexType,
-    WeakIntegralType,
-    _is_weak_dtype,
-    _strong_dtype_num_kind,
-    _weak_type_num_kind,
+    _resolve_one_strong_one_weak_types,
+    _resolve_one_strong_two_weak_types,
 )
 
 
-def _resolve_one_strong_two_weak_types(st_dtype, dtype1, dtype2, dev):
-    "Resolves weak data types per NEP-0050,"
-    "where the second and third arguments are"
-    "permitted to be weak types"
-    if _is_weak_dtype(st_dtype):
-        raise ValueError
-    if _is_weak_dtype(dtype1):
-        if _is_weak_dtype(dtype2):
-            kind_num1 = _weak_type_num_kind(dtype1)
-            kind_num2 = _weak_type_num_kind(dtype2)
-            st_kind_num = _strong_dtype_num_kind(st_dtype)
-
-            if kind_num1 > st_kind_num:
-                if isinstance(dtype1, WeakIntegralType):
-                    ret_dtype1 = dpt.dtype(ti.default_device_int_type(dev))
-                elif isinstance(dtype1, WeakComplexType):
-                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
-                        ret_dtype1 = dpt.complex64
-                    ret_dtype1 = _to_device_supported_dtype(dpt.complex128, dev)
-                else:
-                    ret_dtype1 = _to_device_supported_dtype(dpt.float64, dev)
-            else:
-                ret_dtype1 = st_dtype
-
-            if kind_num2 > st_kind_num:
-                if isinstance(dtype2, WeakIntegralType):
-                    ret_dtype2 = dpt.dtype(ti.default_device_int_type(dev))
-                elif isinstance(dtype2, WeakComplexType):
-                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
-                        ret_dtype2 = dpt.complex64
-                    ret_dtype2 = _to_device_supported_dtype(dpt.complex128, dev)
-                else:
-                    ret_dtype2 = _to_device_supported_dtype(dpt.float64, dev)
-            else:
-                ret_dtype2 = st_dtype
-
-            return ret_dtype1, ret_dtype2
-
-        max_dt_num_kind, max_dtype = max(
-            [
-                (_strong_dtype_num_kind(st_dtype), st_dtype),
-                (_strong_dtype_num_kind(dtype2), dtype2),
-            ]
-        )
-        dt1_kind_num = _weak_type_num_kind(dtype1)
-        if dt1_kind_num > max_dt_num_kind:
-            if isinstance(dtype1, WeakIntegralType):
-                return dpt.dtype(ti.default_device_int_type(dev)), dtype2
-            if isinstance(dtype1, WeakComplexType):
-                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
-                    return dpt.complex64, dtype2
-                return (
-                    _to_device_supported_dtype(dpt.complex128, dev),
-                    dtype2,
-                )
-            return _to_device_supported_dtype(dpt.float64, dev), dtype2
-        else:
-            return max_dtype, dtype2
-    elif _is_weak_dtype(dtype2):
-        max_dt_num_kind, max_dtype = max(
-            [
-                (_strong_dtype_num_kind(st_dtype), st_dtype),
-                (_strong_dtype_num_kind(dtype1), dtype1),
-            ]
-        )
-        dt2_kind_num = _weak_type_num_kind(dtype2)
-        if dt2_kind_num > max_dt_num_kind:
-            if isinstance(dtype2, WeakIntegralType):
-                return dtype1, dpt.dtype(ti.default_device_int_type(dev))
-            if isinstance(dtype2, WeakComplexType):
-                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
-                    return dtype1, dpt.complex64
-                return (
-                    dtype1,
-                    _to_device_supported_dtype(dpt.complex128, dev),
-                )
-            return dtype1, _to_device_supported_dtype(dpt.float64, dev)
-        else:
-            return dtype1, max_dtype
-    else:
-        # both are strong dtypes
-        # return unmodified
-        return dtype1, dtype2
-
-
-def _resolve_one_strong_one_weak_types(st_dtype, dtype, dev):
-    "Resolves one weak data type with one strong data type per NEP-0050"
-    if _is_weak_dtype(st_dtype):
-        raise ValueError
-    if _is_weak_dtype(dtype):
-        st_kind_num = _strong_dtype_num_kind(st_dtype)
-        kind_num = _weak_type_num_kind(dtype)
-        if kind_num > st_kind_num:
-            if isinstance(dtype, WeakIntegralType):
-                return dpt.dtype(ti.default_device_int_type(dev))
-            if isinstance(dtype, WeakComplexType):
-                if st_dtype is dpt.float16 or st_dtype is dpt.float32:
-                    return dpt.complex64
-                return _to_device_supported_dtype(dpt.complex128, dev)
-            return _to_device_supported_dtype(dpt.float64, dev)
-        else:
-            return st_dtype
-    else:
-        return dtype
-
-
 def _check_clip_dtypes(res_dtype, arg1_dtype, arg2_dtype, sycl_dev):
     "Checks if both types `arg1_dtype` and `arg2_dtype` can be"
     "cast to `res_dtype` according to the rule `safe`"

dpctl/tensor/_type_utils.py

@@ -445,6 +445,112 @@ def _resolve_weak_types_comparisons(o1_dtype, o2_dtype, dev):
         return o1_dtype, o2_dtype
 
 
+def _resolve_one_strong_two_weak_types(st_dtype, dtype1, dtype2, dev):
+    "Resolves weak data types per NEP-0050,"
+    "where the second and third arguments are"
+    "permitted to be weak types"
+    if _is_weak_dtype(st_dtype):
+        raise ValueError
+    if _is_weak_dtype(dtype1):
+        if _is_weak_dtype(dtype2):
+            kind_num1 = _weak_type_num_kind(dtype1)
+            kind_num2 = _weak_type_num_kind(dtype2)
+            st_kind_num = _strong_dtype_num_kind(st_dtype)
+
+            if kind_num1 > st_kind_num:
+                if isinstance(dtype1, WeakIntegralType):
+                    ret_dtype1 = dpt.dtype(ti.default_device_int_type(dev))
+                elif isinstance(dtype1, WeakComplexType):
+                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
+                        ret_dtype1 = dpt.complex64
+                    ret_dtype1 = _to_device_supported_dtype(dpt.complex128, dev)
+                else:
+                    ret_dtype1 = _to_device_supported_dtype(dpt.float64, dev)
+            else:
+                ret_dtype1 = st_dtype
+
+            if kind_num2 > st_kind_num:
+                if isinstance(dtype2, WeakIntegralType):
+                    ret_dtype2 = dpt.dtype(ti.default_device_int_type(dev))
+                elif isinstance(dtype2, WeakComplexType):
+                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
+                        ret_dtype2 = dpt.complex64
+                    ret_dtype2 = _to_device_supported_dtype(dpt.complex128, dev)
+                else:
+                    ret_dtype2 = _to_device_supported_dtype(dpt.float64, dev)
+            else:
+                ret_dtype2 = st_dtype
+
+            return ret_dtype1, ret_dtype2
+
+        max_dt_num_kind, max_dtype = max(
+            [
+                (_strong_dtype_num_kind(st_dtype), st_dtype),
+                (_strong_dtype_num_kind(dtype2), dtype2),
+            ]
+        )
+        dt1_kind_num = _weak_type_num_kind(dtype1)
+        if dt1_kind_num > max_dt_num_kind:
+            if isinstance(dtype1, WeakIntegralType):
+                return dpt.dtype(ti.default_device_int_type(dev)), dtype2
+            if isinstance(dtype1, WeakComplexType):
+                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
+                    return dpt.complex64, dtype2
+                return (
+                    _to_device_supported_dtype(dpt.complex128, dev),
+                    dtype2,
+                )
+            return _to_device_supported_dtype(dpt.float64, dev), dtype2
+        else:
+            return max_dtype, dtype2
+    elif _is_weak_dtype(dtype2):
+        max_dt_num_kind, max_dtype = max(
+            [
+                (_strong_dtype_num_kind(st_dtype), st_dtype),
+                (_strong_dtype_num_kind(dtype1), dtype1),
+            ]
+        )
+        dt2_kind_num = _weak_type_num_kind(dtype2)
+        if dt2_kind_num > max_dt_num_kind:
+            if isinstance(dtype2, WeakIntegralType):
+                return dtype1, dpt.dtype(ti.default_device_int_type(dev))
+            if isinstance(dtype2, WeakComplexType):
+                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
+                    return dtype1, dpt.complex64
+                return (
+                    dtype1,
+                    _to_device_supported_dtype(dpt.complex128, dev),
+                )
+            return dtype1, _to_device_supported_dtype(dpt.float64, dev)
+        else:
+            return dtype1, max_dtype
+    else:
+        # both are strong dtypes
+        # return unmodified
+        return dtype1, dtype2
+
+
+def _resolve_one_strong_one_weak_types(st_dtype, dtype, dev):
+    "Resolves one weak data type with one strong data type per NEP-0050"
+    if _is_weak_dtype(st_dtype):
+        raise ValueError
+    if _is_weak_dtype(dtype):
+        st_kind_num = _strong_dtype_num_kind(st_dtype)
+        kind_num = _weak_type_num_kind(dtype)
+        if kind_num > st_kind_num:
+            if isinstance(dtype, WeakIntegralType):
+                return dpt.dtype(ti.default_device_int_type(dev))
+            if isinstance(dtype, WeakComplexType):
+                if st_dtype is dpt.float16 or st_dtype is dpt.float32:
+                    return dpt.complex64
+                return _to_device_supported_dtype(dpt.complex128, dev)
+            return _to_device_supported_dtype(dpt.float64, dev)
+        else:
+            return st_dtype
+    else:
+        return dtype
+
+
 class finfo_object:
     """
     `numpy.finfo` subclass which returns Python floating-point scalars for
@@ -833,6 +939,8 @@ def _default_accumulation_dtype_fp_types(inp_dt, q):
     "_acceptance_fn_divide",
     "_acceptance_fn_negative",
     "_acceptance_fn_subtract",
+    "_resolve_one_strong_one_weak_types",
+    "_resolve_one_strong_two_weak_types",
     "_resolve_weak_types",
     "_resolve_weak_types_comparisons",
     "_weak_type_num_kind",

dpctl/tensor/_utility_functions.py

@@ -20,10 +20,6 @@
 import dpctl.tensor._tensor_impl as ti
 import dpctl.tensor._tensor_reductions_impl as tri
 import dpctl.utils as du
-from dpctl.tensor._clip import (
-    _resolve_one_strong_one_weak_types,
-    _resolve_one_strong_two_weak_types,
-)
 from dpctl.tensor._elementwise_common import (
     _get_dtype,
     _get_queue_usm_type,
@@ -32,6 +28,10 @@
 )
 
 from ._numpy_helper import normalize_axis_index, normalize_axis_tuple
+from ._type_utils import (
+    _resolve_one_strong_one_weak_types,
+    _resolve_one_strong_two_weak_types,
+)
 
 
 def _boolean_reduction(x, axis, keepdims, func):
@@ -159,6 +159,8 @@ def any(x, /, *, axis=None, keepdims=False):
 
 
 def _validate_diff_shape(sh1, sh2, axis):
+    """Utility for validating that two shapes `sh1` and `sh2`
+    are possible to concatenate along `axis`."""
     if not sh2:
         # scalars will always be accepted
         return True
@@ -173,6 +175,9 @@ def _validate_diff_shape(sh1, sh2, axis):
 
 
 def _concat_diff_input(arr, axis, prepend, append):
+    """Concatenates `arr`, `prepend` and, `append` along `axis`,
+    where `arr` is an array and `prepend` and `append` are
+    any mixture of arrays and scalars."""
     if prepend is not None and append is not None:
         q1, x_usm_type = arr.sycl_queue, arr.usm_type
         q2, prepend_usm_type = _get_queue_usm_type(prepend)