UnaryElementwiseFunc class now takes an acceptance function

This change was made to mirror promotion behavior of divide in reciprocal

Adds getter method for the acceptance function

Adds tests for reciprocal

Author: ndgrigorian

dpctl/tensor/_elementwise_common.py

@@ -28,7 +28,8 @@
 
 from ._copy_utils import _empty_like_orderK, _empty_like_pair_orderK
 from ._type_utils import (
-    _acceptance_fn_default,
+    _acceptance_fn_default1,
+    _acceptance_fn_default2,
     _all_data_types,
     _find_buf_dtype,
     _find_buf_dtype2,
@@ -62,17 +63,39 @@ class UnaryElementwiseFunc:
             computational tasks complete execution, while the second event
             corresponds to computational tasks associated with function
             evaluation.
+        acceptance_fn (callable, optional):
+            Function to influence type promotion behavior of this unary
+            function. The function takes 4 arguments:
+                arg_dtype - Data type of the first argument
+                buf_dtype - Data type the argument would be cast to
+                res_dtype - Data type of the output array with function values
+                sycl_dev - The :class:`dpctl.SyclDevice` where the function
+                    evaluation is carried out.
+            The function is invoked when the argument of the unary function
+            requires casting, e.g. the argument of `dpctl.tensor.log` is an
+            array with integral data type.
         docs (str):
             Documentation string for the unary function.
     """
 
-    def __init__(self, name, result_type_resolver_fn, unary_dp_impl_fn, docs):
+    def __init__(
+        self,
+        name,
+        result_type_resolver_fn,
+        unary_dp_impl_fn,
+        docs,
+        acceptance_fn=None,
+    ):
         self.__name__ = "UnaryElementwiseFunc"
         self.name_ = name
         self.result_type_resolver_fn_ = result_type_resolver_fn
         self.types_ = None
         self.unary_fn_ = unary_dp_impl_fn
         self.__doc__ = docs
+        if callable(acceptance_fn):
+            self.acceptance_fn_ = acceptance_fn
+        else:
+            self.acceptance_fn_ = _acceptance_fn_default1
 
     def __str__(self):
         return f"<{self.__name__} '{self.name_}'>"
@@ -93,6 +116,24 @@ def get_type_result_resolver_function(self):
         """
         return self.result_type_resolver_fn_
 
+    def get_type_promotion_path_acceptance_function(self):
+        """Returns the acceptance function for this
+        elementwise binary function.
+
+        Acceptance function influences the type promotion
+        behavior of this unary function.
+        The function takes 4 arguments:
+            arg_dtype - Data type of the first argument
+            buf_dtype - Data type the argument would be cast to
+            res_dtype - Data type of the output array with function values
+            sycl_dev - The :class:`dpctl.SyclDevice` where the function
+                evaluation is carried out.
+        The function is invoked when the argument of the unary function
+        requires casting, e.g. the argument of `dpctl.tensor.log` is an
+        array with integral data type.
+        """
+        return self.acceptance_fn_
+
     @property
     def types(self):
         """Returns information about types supported by
@@ -122,7 +163,10 @@ def __call__(self, x, out=None, order="K"):
         if order not in ["C", "F", "K", "A"]:
             order = "K"
         buf_dt, res_dt = _find_buf_dtype(
-            x.dtype, self.result_type_resolver_fn_, x.sycl_device
+            x.dtype,
+            self.result_type_resolver_fn_,
+            x.sycl_device,
+            acceptance_fn=self.acceptance_fn_,
         )
         if res_dt is None:
             raise TypeError(
@@ -482,7 +526,7 @@ def __init__(
         if callable(acceptance_fn):
             self.acceptance_fn_ = acceptance_fn
         else:
-            self.acceptance_fn_ = _acceptance_fn_default
+            self.acceptance_fn_ = _acceptance_fn_default2
 
     def __str__(self):
         return f"<{self.__name__} '{self.name_}'>"

dpctl/tensor/_elementwise_funcs.py

@@ -17,7 +17,7 @@
 import dpctl.tensor._tensor_elementwise_impl as ti
 
 from ._elementwise_common import BinaryElementwiseFunc, UnaryElementwiseFunc
-from ._type_utils import _acceptance_fn_divide
+from ._type_utils import _acceptance_fn_divide, _acceptance_fn_reciprocal
 
 # U01: ==== ABS    (x)
 _abs_docstring_ = """
@@ -1916,6 +1916,7 @@
     ti._reciprocal_result_type,
     ti._reciprocal,
     _reciprocal_docstring,
+    acceptance_fn=_acceptance_fn_reciprocal,
 )
 
 

dpctl/tensor/_type_utils.py

@@ -132,7 +132,27 @@ def _to_device_supported_dtype(dt, dev):
     return dt
 
 
-def _find_buf_dtype(arg_dtype, query_fn, sycl_dev):
+def _acceptance_fn_default1(arg_dtype, ret_buf_dt, res_dt, sycl_dev):
+    return True
+
+
+def _acceptance_fn_reciprocal(arg_dtype, buf_dt, res_dt, sycl_dev):
+    # if the kind of result is different from
+    # the kind of input, use the default data
+    # we use default dtype for the resulting kind.
+    # This guarantees alignment of reciprocal and
+    # divide output types.
+    if buf_dt.kind != arg_dtype.kind:
+        default_dt = _get_device_default_dtype(res_dt.kind, sycl_dev)
+        if res_dt == default_dt:
+            return True
+        else:
+            return False
+    else:
+        return True
+
+
+def _find_buf_dtype(arg_dtype, query_fn, sycl_dev, acceptance_fn):
     res_dt = query_fn(arg_dtype)
     if res_dt:
         return None, res_dt
@@ -144,7 +164,11 @@ def _find_buf_dtype(arg_dtype, query_fn, sycl_dev):
         if _can_cast(arg_dtype, buf_dt, _fp16, _fp64):
             res_dt = query_fn(buf_dt)
             if res_dt:
-                return buf_dt, res_dt
+                acceptable = acceptance_fn(arg_dtype, buf_dt, res_dt, sycl_dev)
+                if acceptable:
+                    return buf_dt, res_dt
+                else:
+                    continue
 
     return None, None
 
@@ -163,7 +187,7 @@ def _get_device_default_dtype(dt_kind, sycl_dev):
     raise RuntimeError
 
 
-def _acceptance_fn_default(
+def _acceptance_fn_default2(
     arg1_dtype, arg2_dtype, ret_buf1_dt, ret_buf2_dt, res_dt, sycl_dev
 ):
     return True
@@ -230,6 +254,8 @@ def _find_buf_dtype2(arg1_dtype, arg2_dtype, query_fn, sycl_dev, acceptance_fn):
     "_find_buf_dtype",
     "_find_buf_dtype2",
     "_to_device_supported_dtype",
-    "_acceptance_fn_default",
+    "_acceptance_fn_default1",
+    "_acceptance_fn_reciprocal",
+    "_acceptance_fn_default2",
     "_acceptance_fn_divide",
 ]

dpctl/tests/elementwise/test_reciprocal.py

@@ -0,0 +1,93 @@
+#                       Data Parallel Control (dpctl)
+#
+#  Copyright 2020-2023 Intel Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import itertools
+
+import pytest
+
+import dpctl.tensor as dpt
+from dpctl.tests.helper import get_queue_or_skip, skip_if_dtype_not_supported
+
+from .utils import _all_dtypes, _complex_fp_dtypes
+
+
+@pytest.mark.parametrize("dtype", _all_dtypes)
+def test_reciprocal_out_type(dtype):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dtype, q)
+
+    x = dpt.asarray(1, dtype=dtype, sycl_queue=q)
+    one = dpt.asarray(1, dtype=dtype, sycl_queue=q)
+    expected_dtype = dpt.divide(one, x).dtype
+    assert dpt.reciprocal(x).dtype == expected_dtype
+
+
+@pytest.mark.parametrize("dtype", _all_dtypes)
+def test_reciprocal_output_contig(dtype):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dtype, q)
+
+    n_seq = 1027
+
+    x = dpt.linspace(1, 13, num=n_seq, dtype=dtype, sycl_queue=q)
+    res = dpt.reciprocal(x)
+    expected = 1 / x
+    tol = 8 * dpt.finfo(res.dtype).resolution
+    assert dpt.allclose(res, expected, atol=tol, rtol=tol)
+
+
+@pytest.mark.parametrize("dtype", _all_dtypes)
+def test_reciprocal_output_strided(dtype):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dtype, q)
+
+    n_seq = 2054
+
+    x = dpt.linspace(1, 13, num=n_seq, dtype=dtype, sycl_queue=q)[::-2]
+    res = dpt.reciprocal(x)
+    expected = 1 / x
+    tol = 8 * dpt.finfo(res.dtype).resolution
+    assert dpt.allclose(res, expected, atol=tol, rtol=tol)
+
+
+def test_reciprocal_special_cases():
+    get_queue_or_skip()
+
+    x = dpt.asarray([dpt.nan, 0.0, -0.0, dpt.inf, -dpt.inf], dtype="f4")
+    res = dpt.reciprocal(x)
+    expected = dpt.asarray([dpt.nan, dpt.inf, -dpt.inf, 0.0, -0.0], dtype="f4")
+    assert dpt.allclose(res, expected, equal_nan=True)
+
+
+@pytest.mark.parametrize("dtype", _complex_fp_dtypes)
+def test_reciprocal_complex_special_cases(dtype):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dtype, q)
+
+    nans_ = [dpt.nan, -dpt.nan]
+    infs_ = [dpt.inf, -dpt.inf]
+    finites_ = [-1.0, -0.0, 0.0, 1.0]
+    inps_ = nans_ + infs_ + finites_
+    c_ = [complex(*v) for v in itertools.product(inps_, repeat=2)]
+
+    z = dpt.asarray(c_, dtype=dtype)
+    r = dpt.reciprocal(z)
+
+    expected = 1 / z
+
+    tol = dpt.finfo(r.dtype).resolution
+
+    assert dpt.allclose(r, expected, atol=tol, rtol=tol, equal_nan=True)

dpctl/tests/elementwise/test_type_utils.py

@@ -124,7 +124,9 @@ def _denier_fn(dt):
         for fp16 in [True, False]:
             dev = MockDevice(fp16, fp64)
             arg_dt = dpt.float64
-            r = tu._find_buf_dtype(arg_dt, _denier_fn, dev)
+            r = tu._find_buf_dtype(
+                arg_dt, _denier_fn, dev, tu._acceptance_fn_default1
+            )
             assert r == (
                 None,
                 None,
@@ -157,7 +159,7 @@ def _denier_fn(dt1, dt2):
             arg1_dt = dpt.float64
             arg2_dt = dpt.complex64
             r = tu._find_buf_dtype2(
-                arg1_dt, arg2_dt, _denier_fn, dev, tu._acceptance_fn_default
+                arg1_dt, arg2_dt, _denier_fn, dev, tu._acceptance_fn_default2
             )
             assert r == (
                 None,