Implements logsumexp and reduce_hypot

Author: ndgrigorian

dpctl/tensor/__init__.py

@@ -164,7 +164,16 @@
     tanh,
     trunc,
 )
-from ._reduction import argmax, argmin, max, min, prod, sum
+from ._reduction import (
+    argmax,
+    argmin,
+    logsumexp,
+    max,
+    min,
+    prod,
+    reduce_hypot,
+    sum,
+)
 from ._testing import allclose
 
 __all__ = [
@@ -322,4 +331,6 @@
     "exp2",
     "copysign",
     "rsqrt",
+    "logsumexp",
+    "reduce_hypot",
 ]

dpctl/tensor/_reduction.py

@@ -52,6 +52,28 @@ def _default_reduction_dtype(inp_dt, q):
     return res_dt
 
 
+def _default_reduction_dtype_fp_types(inp_dt, q):
+    """Gives default output data type for given input data
+    type `inp_dt` when reduction is performed on queue `q`
+    and the reduction supports only floating-point data types
+    """
+    inp_kind = inp_dt.kind
+    if inp_kind in "biu":
+        res_dt = dpt.dtype(ti.default_device_fp_type(q))
+        can_cast_v = dpt.can_cast(inp_dt, res_dt)
+        if not can_cast_v:
+            _fp64 = q.sycl_device.has_aspect_fp64
+            res_dt = dpt.float64 if _fp64 else dpt.float32
+    elif inp_kind in "f":
+        res_dt = dpt.dtype(ti.default_device_fp_type(q))
+        if res_dt.itemsize < inp_dt.itemsize:
+            res_dt = inp_dt
+    elif inp_kind in "c":
+        raise TypeError("reduction not defined for complex types")
+
+    return res_dt
+
+
 def _reduction_over_axis(
     x,
     axis,
@@ -91,12 +113,15 @@ def _reduction_over_axis(
                 res_shape = res_shape + (1,) * red_nd
                 inv_perm = sorted(range(nd), key=lambda d: perm[d])
                 res_shape = tuple(res_shape[i] for i in inv_perm)
-            return dpt.full(
-                res_shape,
-                _identity,
-                dtype=res_dt,
-                usm_type=res_usm_type,
-                sycl_queue=q,
+            return dpt.astype(
+                dpt.full(
+                    res_shape,
+                    _identity,
+                    dtype=_default_reduction_type_fn(inp_dt, q),
+                    usm_type=res_usm_type,
+                    sycl_queue=q,
+                ),
+                res_dt,
             )
     if red_nd == 0:
         return dpt.astype(x, res_dt, copy=False)
@@ -116,7 +141,7 @@ def _reduction_over_axis(
                 "Automatically determined reduction data type does not "
                 "have direct implementation"
             )
-        tmp_dt = _default_reduction_dtype(inp_dt, q)
+        tmp_dt = _default_reduction_type_fn(inp_dt, q)
         tmp = dpt.empty(
             res_shape, dtype=tmp_dt, usm_type=res_usm_type, sycl_queue=q
         )
@@ -161,13 +186,13 @@ def sum(x, axis=None, dtype=None, keepdims=False):
                   the returned array will have the default real-valued
                   floating-point data type for the device where input
                   array `x` is allocated.
-                * If x` has signed integral data type, the returned array
+                * If `x` has signed integral data type, the returned array
                   will have the default signed integral type for the device
                   where input array `x` is allocated.
                 * If `x` has unsigned integral data type, the returned array
                   will have the default unsigned integral type for the device
                   where input array `x` is allocated.
-                * If `x` has a complex-valued floating-point data typee,
+                * If `x` has a complex-valued floating-point data type,
                   the returned array will have the default complex-valued
                   floating-pointer data type for the device where input
                   array `x` is allocated.
@@ -222,13 +247,13 @@ def prod(x, axis=None, dtype=None, keepdims=False):
                   the returned array will have the default real-valued
                   floating-point data type for the device where input
                   array `x` is allocated.
-                * If x` has signed integral data type, the returned array
+                * If `x` has signed integral data type, the returned array
                   will have the default signed integral type for the device
                   where input array `x` is allocated.
                 * If `x` has unsigned integral data type, the returned array
                   will have the default unsigned integral type for the device
                   where input array `x` is allocated.
-                * If `x` has a complex-valued floating-point data typee,
+                * If `x` has a complex-valued floating-point data type,
                   the returned array will have the default complex-valued
                   floating-pointer data type for the device where input
                   array `x` is allocated.
@@ -263,6 +288,114 @@ def prod(x, axis=None, dtype=None, keepdims=False):
     )
 
 
+def logsumexp(x, axis=None, dtype=None, keepdims=False):
+    """logsumexp(x, axis=None, dtype=None, keepdims=False)
+
+    Calculates the logarithm of the sum of exponentials of elements in the
+    input array `x`.
+
+    Args:
+        x (usm_ndarray):
+            input array.
+        axis (Optional[int, Tuple[int, ...]]):
+            axis or axes along which values must be computed. If a tuple
+            of unique integers, values are computed over multiple axes.
+            If `None`, the result is computed over the entire array.
+            Default: `None`.
+        dtype (Optional[dtype]):
+            data type of the returned array. If `None`, the default data
+            type is inferred from the "kind" of the input array data type.
+                * If `x` has a real-valued floating-point data type,
+                  the returned array will have the default real-valued
+                  floating-point data type for the device where input
+                  array `x` is allocated.
+                * If `x` has a boolean or integral data type, the returned array
+                  will have the default floating point data type for the device
+                  where input array `x` is allocated.
+                * If `x` has a complex-valued floating-point data type,
+                  an error is raised.
+            If the data type (either specified or resolved) differs from the
+            data type of `x`, the input array elements are cast to the
+            specified data type before computing the result. Default: `None`.
+        keepdims (Optional[bool]):
+            if `True`, the reduced axes (dimensions) are included in the result
+            as singleton dimensions, so that the returned array remains
+            compatible with the input arrays according to Array Broadcasting
+            rules. Otherwise, if `False`, the reduced axes are not included in
+            the returned array. Default: `False`.
+    Returns:
+        usm_ndarray:
+            an array containing the results. If the result was computed over
+            the entire array, a zero-dimensional array is returned. The returned
+            array has the data type as described in the `dtype` parameter
+            description above.
+    """
+    return _reduction_over_axis(
+        x,
+        axis,
+        dtype,
+        keepdims,
+        ti._logsumexp_over_axis,
+        ti._logsumexp_over_axis_dtype_supported,
+        _default_reduction_dtype_fp_types,
+        _identity=-dpt.inf,
+    )
+
+
+def reduce_hypot(x, axis=None, dtype=None, keepdims=False):
+    """reduce_hypot(x, axis=None, dtype=None, keepdims=False)
+
+    Calculates the square root of the sum of squares of elements in the input
+    array `x`.
+
+    Args:
+        x (usm_ndarray):
+            input array.
+        axis (Optional[int, Tuple[int, ...]]):
+            axis or axes along which values must be computed. If a tuple
+            of unique integers, values are computed over multiple axes.
+            If `None`, the result is computed over the entire array.
+            Default: `None`.
+        dtype (Optional[dtype]):
+            data type of the returned array. If `None`, the default data
+            type is inferred from the "kind" of the input array data type.
+                * If `x` has a real-valued floating-point data type,
+                  the returned array will have the default real-valued
+                  floating-point data type for the device where input
+                  array `x` is allocated.
+                * If `x` has a boolean or integral data type, the returned array
+                  will have the default floating point data type for the device
+                  where input array `x` is allocated.
+                * If `x` has a complex-valued floating-point data type,
+                  an error is raised.
+            If the data type (either specified or resolved) differs from the
+            data type of `x`, the input array elements are cast to the
+            specified data type before computing the result. Default: `None`.
+        keepdims (Optional[bool]):
+            if `True`, the reduced axes (dimensions) are included in the result
+            as singleton dimensions, so that the returned array remains
+            compatible with the input arrays according to Array Broadcasting
+            rules. Otherwise, if `False`, the reduced axes are not included in
+            the returned array. Default: `False`.
+    Returns:
+        usm_ndarray:
+            an array containing the results. If the result was computed over
+            the entire array, a zero-dimensional array is returned. The returned
+            array has the data type as described in the `dtype` parameter
+            description above.
+    """
+    return _reduction_over_axis(
+        x,
+        axis,
+        dtype,
+        keepdims,
+        ti._hypot_over_axis,
+        ti._hypot_over_axis_dtype_supported,
+        _default_reduction_dtype_fp_types,
+        _identity=0,
+    )
+
+
 def _comparison_over_axis(x, axis, keepdims, _reduction_fn):
     if not isinstance(x, dpt.usm_ndarray):
         raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x)}")