!31810 [assistant][ops] Add new LogAddExp2

Merge pull request !31810 from wxy220/LogAddExp2

mindspore/python/mindspore/ops/function/__init__.py

@@ -31,7 +31,7 @@ from .math_func import (addn, absolute, abs, tensor_add, add, neg_tensor, neg, t
                         lp_norm, round, tensor_gt, gt, tensor_ge, ge, tensor_sub, sub, tensor_mul, mul, tensor_div, div,
                         tensor_floordiv, floor_div, floordiv, tensor_pow, pow, pows, tensor_mod, floor_mod, floormod,
                         tensor_exp, exp, tensor_expm1, expm1, equal, not_equal, ne, isfinite, isnan, same_type_shape,
-                        log, log_matrix_determinant, matrix_determinant, maximum,
+                        log, log_matrix_determinant, matrix_determinant, maximum, logaddexp2,
                         invert, minimum, floor, logical_not, logical_or, logical_and, sin, cos, tan,
                         asin, acos, atan, sinh, cosh, tanh, asinh, acosh, atanh, atan2, bitwise_and, bitwise_or,
                         bitwise_xor, erf, erfc, cdist, bessel_i0, bessel_i0e, bessel_j0, bessel_j1, bessel_k0,

mindspore/python/mindspore/ops/function/math_func.py

@@ -15,9 +15,19 @@
 
 """Defines math operators with functional form."""
 
+from mindspore.ops.primitive import constexpr
 from mindspore.ops import operations as P
 from ..operations.math_ops import (BesselJ0, BesselJ1, BesselK0, BesselK0e, BesselY0, BesselY1, BesselK1,
                                    BesselK1e)
+from ...common.tensor import Tensor
+from ..._c_expression import Tensor as Tensor_
+
+
+@constexpr
+def _make_tensor(val, dtype):
+    """Returns the tensor with value `val` and dtype `dtype`."""
+    return Tensor(val, dtype)
+
 
 #####################################
 # Public Operation Functions.
@@ -2172,6 +2182,52 @@ def minimum(x, y):
     return minimum_(x, y)
 
 
+def logaddexp2(x1, x2):
+    """
+    Computes the logarithm of the sum of exponentiations in base of 2 of the inputs.
+
+    Calculates ``log2(2**x1 + 2**x2)``. This function is useful in machine learning when the computed
+    probability of an event may be small beyond the range of normal floating point numbers.
+    In this case, the base-2 logarithm of the calculated probability can be used instead.
+    This function allows to add probabilities stored in this way.
+
+    Args:
+        x1 (Tensor): Input tensor.
+        x2 (Tensor): Input tensor. If ``x1.shape != x2.shape``, they must be broadcastable to
+            a common shape (which becomes the shape of the output).
+
+    Returns:
+        Tensor or scalar. This is a scalar if both `x1` and `x2` are scalars.
+
+    Raises:
+        TypeError: If `x1`, `x2` is not a Tensor.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> x1 = Tensor(np.array([2, 4, 8]).astype(np.float16))
+        >>> x2 = Tensor(np.array([2]).astype(np.float16))
+        >>> output = ops.logaddexp2(x1, x2)
+        >>> print(output)
+        [3. 4.32 8.02]
+    """
+
+    log_op = P.Log()
+    pow_op = P.Pow()
+    add_op = P.Add()
+
+    if not isinstance(x1, (Tensor, Tensor_)):
+        raise TypeError("The input x1 must be Tensor.")
+    if not isinstance(x2, (Tensor, Tensor_)):
+        raise TypeError("The input x2 must be Tensor.")
+
+    add_exp = add_op(pow_op(2, x1), pow_op(2, x2))
+    tensor_2 = _make_tensor(2, add_exp.dtype)
+
+    return log_op(add_exp) / log_op(tensor_2)
+
+
 def cdist(x, y, p=2.0):
     """
     Computes batched the p-norm distance between each pair of the two collections of row vectors.
@@ -2400,6 +2456,7 @@ __all__ = [
     'neg',
     'tensor_lt',
     'less',
+    'logaddexp2',
     'tensor_le',
     'le',
     'lerp',

tests/ut/python/ops/test_math_ops.py

@@ -22,6 +22,7 @@ import mindspore.context as context
 import mindspore.nn as nn
 from mindspore import Tensor
 from mindspore.common import dtype as mstype
+from mindspore import ops
 from mindspore.ops import composite as C
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
@@ -386,6 +387,16 @@ class ErfcNet(nn.Cell):
         return self.erfc(x)
 
 
+class LogAddExp2Func(nn.Cell):
+    def __init__(self):
+        super(LogAddExp2Func, self).__init__()
+        self.logaddexp2 = ops.logaddexp2
+
+    def construct(self, x1, x2):
+        y = self.logaddexp2(x1, x2)
+        return y
+
+
 test_case_math_ops = [
     ('MatMulGrad', {
         'block': GradWrap(NetWithLoss(MatMulNet())),
@@ -441,6 +452,11 @@ test_case_math_ops = [
         'desc_inputs': [Tensor(np.array([[1.0, 2.0, 4.0]], np.float32))],
         'desc_bprop': [Tensor(np.array([[1.0, 2.0, 4.0]], np.float32))],
     }),
+    ('LogAddExp2', {
+        'block': LogAddExp2Func(),
+        'desc_inputs': [Tensor(np.array([1.0, 2.0, 3.0], np.float16)), Tensor(np.array([2.0], np.float16))],
+        'desc_bprop': [Tensor(np.array([1.0, 2.0, 3.0], np.float16)), Tensor(np.array([2.0], np.float16))],
+    }),
 ]
 
 test_case_lists = [test_case_math_ops]