• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020-2021 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""math"""
16from __future__ import absolute_import
17
18import numpy as np
19
20from mindspore import log as logger
21from mindspore.ops import operations as P
22from mindspore.common.tensor import Tensor
23from mindspore.common._decorator import deprecated
24from mindspore.ops.primitive import constexpr, _primexpr
25from mindspore.ops import functional as F
26from mindspore.nn.cell import Cell
27from mindspore.common import dtype as mstype
28from mindspore import _checkparam as validator
29
30__all__ = ['ReduceLogSumExp',
31           'Range',
32           'LGamma',
33           'DiGamma',
34           'IGamma',
35           'LBeta',
36           'CosineSimilarity',
37           'MatMul',
38           'Moments',
39           'MatInverse',
40           'MatDet',
41           ]
42
43_BASE_LANCZOS_COEFF = 0.99999999999980993227684700473478
44_LANCZOS_COEFFICIENTS = [676.520368121885098567009190444019,
45                         -1259.13921672240287047156078755283,
46                         771.3234287776530788486528258894,
47                         -176.61502916214059906584551354,
48                         12.507343278686904814458936853,
49                         -0.13857109526572011689554707,
50                         9.984369578019570859563e-6,
51                         1.50563273514931155834e-7]
52
53
54@constexpr
55def _check_input_dtype(param_name, input_dtype, allow_dtypes, cls_name):
56    validator.check_type_name(param_name, input_dtype, allow_dtypes, cls_name)
57
58
59class ReduceLogSumExp(Cell):
60    r"""
61    Reduces a dimension of a tensor by calculating exponential for all elements in the dimension,
62    then calculate logarithm of the sum.
63
64    .. math::
65
66        ReduceLogSumExp(x) = \log(\sum(e^x))
67
68    Args:
69        axis (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: ``()`` , reduce all dimensions.
70            Only constant value is allowed.
71        keep_dims (bool): If True, keep these reduced dimensions and the length is 1.
72            If False, don't keep these dimensions.
73            Default : False.
74
75    Inputs:
76        - **x** (Tensor) - The input tensor. With float16 or float32 data type.
77
78    Outputs:
79        Tensor, has the same dtype as the `x`.
80
81        - If axis is (), and keep_dims is False,
82          the output is a 0-D tensor representing the sum of all elements in the input tensor.
83        - If axis is int, set as 2, and keep_dims is False,
84          the shape of output is :math:`(x_1, x_3, ..., x_R)`.
85        - If axis is tuple(int), set as (2, 3), and keep_dims is False,
86          the shape of output is :math:`(x_1, x_4, ..., x_R)`.
87
88    Raises:
89        TypeError: If `axis` is not one of int, list, tuple.
90        TypeError: If `keep_dims` is not bool.
91        TypeError: If dtype of `x` is neither float16 nor float32.
92
93    Supported Platforms:
94        ``Ascend`` ``GPU`` ``CPU``
95
96    Examples:
97        >>> import mindspore as ms
98        >>> import numpy as np
99        >>> x = ms.Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
100        >>> op = ms.nn.ReduceLogSumExp(1, keep_dims=True)
101        >>> output = op(x)
102        >>> print(output.shape)
103        (3, 1, 5, 6)
104    """
105
106    def __init__(self, axis, keep_dims=False):
107        """Initialize ReduceLogSumExp."""
108        super(ReduceLogSumExp, self).__init__()
109        validator.check_value_type('axis', axis, [int, list, tuple], self.cls_name)
110        validator.check_value_type('keep_dims', keep_dims, [bool], self.cls_name)
111        self.axis = axis
112        self.exp = P.Exp()
113        self.sum = P.ReduceSum(keep_dims)
114        self.log = P.Log()
115        self.max = P.ReduceMax()
116
117    def construct(self, x):
118        x_max = self.max(x)
119        exp = self.exp(x - x_max)
120        sumexp = self.sum(exp, self.axis)
121        logsumexp = self.log(sumexp)
122        return logsumexp + x_max
123
124
125class Range(Cell):
126    r"""
127    The Range class will be deprecated in the future,
128    this function can be replaced by :func:`ops.range`
129    """
130
131    def __init__(self, start, limit=None, delta=1):
132        """Initialize Range."""
133        super(Range, self).__init__()
134        logger.warning("The Range class will be deprecated in the future,"
135                       "this function can be replaced by :func:`ops.range`")
136        if delta == 0:
137            raise ValueError(f"For '{self.cls_name}', the 'delta' can not be zero.")
138        data = np.arange(start, limit, delta)
139        if data.dtype == np.float_:
140            self.ms_dtype = mstype.float32
141        else:
142            self.ms_dtype = mstype.int32
143        self.result_tensor = Tensor(data, dtype=self.ms_dtype)
144
145    def construct(self):
146        return self.result_tensor
147
148
149class LGamma(Cell):
150    r"""
151    Calculates LGamma using Lanczos' approximation referring to "A Precision Approximation of the Gamma Function".
152    The algorithm is:
153
154    .. math::
155        \begin{array}{ll} \\
156            lgamma(z + 1) = \frac{(\log(2) + \log(pi))}{2} + (z + 1/2) * log(t(z)) - t(z) + A(z) \\
157            t(z) = z + kLanczosGamma + 1/2 \\
158            A(z) = kBaseLanczosCoeff + \sum_{k=1}^n \frac{kLanczosCoefficients[i]}{z + k}
159        \end{array}
160
161    However, if the input is less than 0.5 use Euler's reflection formula:
162
163    .. math::
164
165        lgamma(x) = \log(pi) - lgamma(1-x) - \log(abs(sin(pi * x)))
166
167    And please note that
168
169    .. math::
170
171        lgamma(+/-inf) = +inf
172
173    Thus, the behaviour of LGamma follows:
174
175    - when x > 0.5, return log(Gamma(x))
176    - when x < 0.5 and is not an integer, return the real part of Log(Gamma(x)) where Log is the complex logarithm
177    - when x is an integer less or equal to 0, return +inf
178    - when x = +/- inf, return +inf
179
180    Inputs:
181        - **x** (Tensor) - The input tensor. Only float16, float32 are supported.
182
183    Outputs:
184        Tensor, has the same shape and dtype as the `x`.
185
186    Raises:
187        TypeError: If dtype of `x` is neither float16 nor float32.
188
189    Supported Platforms:
190        ``Ascend`` ``GPU``
191
192    Examples:
193        >>> import mindspore as ms
194        >>> import numpy as np
195        >>> x = ms.Tensor(np.array([2, 3, 4]).astype(np.float32))
196        >>> op = ms.nn.LGamma()
197        >>> output = op(x)
198        >>> print(output)
199        [3.5762787e-07 6.9314754e-01 1.7917603e+00]
200    """
201
202    def __init__(self):
203        """Initialize LGamma."""
204        super(LGamma, self).__init__()
205        # const numbers
206        self.k_lanczos_gamma = 7
207        self.k_base_lanczos_coeff = _BASE_LANCZOS_COEFF
208        self.k_lanczos_coefficients = _LANCZOS_COEFFICIENTS
209        self.one_half = 0.5
210        self.one = 1
211        self.two = 2
212        self.inf = np.inf
213        self.pi = np.pi
214        self.log_2 = np.log(self.two)
215        self.log_pi = np.log(np.pi)
216        self.log_sqrt_two_pi = (self.log_2 + self.log_pi) / self.two
217        self.lanczos_gamma_plus_one_half = self.k_lanczos_gamma + 0.5
218        self.log_lanczos_gamma_plus_one_half = np.log(self.lanczos_gamma_plus_one_half)
219
220        # operations
221        self.log = P.Log()
222        self.log1p = P.Log1p()
223        self.abs = P.Abs()
224        self.shape = P.Shape()
225        self.dtype = P.DType()
226        self.floor = P.Floor()
227        self.equal = P.Equal()
228        self.greater = P.Greater()
229        self.less = P.Less()
230        self.lessequal = P.LessEqual()
231        self.select = P.Select()
232        self.sin = P.Sin()
233        self.isfinite = P.IsFinite()
234        self.ones_like = P.OnesLike()
235
236    def construct(self, x):
237        input_dtype = self.dtype(x)
238        _check_input_dtype("x", input_dtype, [mstype.float16, mstype.float32], self.cls_name)
239        if F.is_sequence_value_unknown(self.shape(x)):
240            infinity = self.ones_like(x) * F.cast(self.inf, input_dtype)
241        else:
242            infinity = F.fill(input_dtype, self.shape(x), self.inf)
243
244        need_to_reflect = self.less(x, 0.5)
245        neg_input = -x
246        z = self.select(need_to_reflect, neg_input, x - 1)
247
248        @constexpr
249        def _calculate_reflected_x(z, k_base_lanczos_coeff, k_lanczos_coefficients):
250            reflex_x = k_base_lanczos_coeff
251            for i in range(8):
252                product_ = k_lanczos_coefficients[i] / (z + i + 1)
253                reflex_x = product_ + reflex_x
254            return reflex_x
255
256        reflex_x = _calculate_reflected_x(z, self.k_base_lanczos_coeff, self.k_lanczos_coefficients)
257
258        t = z + self.lanczos_gamma_plus_one_half
259        log_t = self.log1p(z / self.lanczos_gamma_plus_one_half) + self.log_lanczos_gamma_plus_one_half
260
261        log_y = self.log(reflex_x) + (z + self.one_half - t / log_t) * log_t + self.log_sqrt_two_pi
262
263        abs_input = self.abs(x)
264        abs_frac_input = abs_input - self.floor(abs_input)
265        x = self.select(self.lessequal(x, 0.0), self.select(self.equal(abs_frac_input, 0.0), infinity, x), x)
266        reduced_frac_input = self.select(self.greater(abs_frac_input, 0.5),
267                                         1 - abs_frac_input, abs_frac_input)
268        reflection_denom = self.log(self.sin(self.pi * reduced_frac_input))
269
270        reflection = self.select(self.isfinite(reflection_denom),
271                                 -reflection_denom - log_y + self.log_pi,  # pylint: disable=invalid-unary-operand-type
272                                 -reflection_denom)  # pylint: disable=invalid-unary-operand-type
273
274        result = self.select(need_to_reflect, reflection, log_y)
275
276        return self.select(self.isfinite(x), result, infinity)
277
278
279class DiGamma(Cell):
280    r"""
281    Calculates Digamma using Lanczos' approximation referring to "A Precision Approximation of the Gamma Function".
282    The algorithm is:
283
284    .. math::
285        \begin{array}{ll} \\
286            digamma(z + 1) = log(t(z)) + A'(z) / A(z) - kLanczosGamma / t(z) \\
287            t(z) = z + kLanczosGamma + 1/2 \\
288            A(z) = kBaseLanczosCoeff + \sum_{k=1}^n \frac{kLanczosCoefficients[i]}{z + k} \\
289            A'(z) = \sum_{k=1}^n \frac{kLanczosCoefficients[i]}{{z + k}^2}
290        \end{array}
291
292    However, if the input is less than 0.5 use Euler's reflection formula:
293
294    .. math::
295
296        digamma(x) = digamma(1 - x) - pi * cot(pi * x)
297
298    Inputs:
299        - **x** (Tensor[Number]) - The input tensor. Only float16, float32 are supported.
300
301    Outputs:
302        Tensor, has the same shape and dtype as the `x`.
303
304    Raises:
305        TypeError: If dtype of `x` is neither float16 nor float32.
306
307    Supported Platforms:
308        ``Ascend`` ``GPU``
309
310    Examples:
311        >>> import mindspore as ms
312        >>> import numpy as np
313        >>> x = ms.Tensor(np.array([2, 3, 4]).astype(np.float32))
314        >>> op = ms.nn.DiGamma()
315        >>> output = op(x)
316        >>> print(output)
317        [0.42278463  0.92278427 1.2561178]
318    """
319
320    def __init__(self):
321        """Initialize DiGamma."""
322        super(DiGamma, self).__init__()
323        # const numbers
324        self.k_lanczos_gamma = 7
325        self.k_base_lanczos_coeff = _BASE_LANCZOS_COEFF
326        self.k_lanczos_coefficients = _LANCZOS_COEFFICIENTS
327        self.nan = np.nan
328        self.pi = np.pi
329        self.lanczos_gamma_plus_one_half = self.k_lanczos_gamma + 0.5
330        self.log_lanczos_gamma_plus_one_half = np.log(self.lanczos_gamma_plus_one_half)
331
332        # operations
333        self.log1p = P.Log1p()
334        self.abs = P.Abs()
335        self.shape = P.Shape()
336        self.dtype = P.DType()
337        self.floor = P.Floor()
338        self.equal = P.Equal()
339        self.less = P.Less()
340        self.select = P.Select()
341        self.sin = P.Sin()
342        self.cos = P.Cos()
343        self.logicaland = P.LogicalAnd()
344
345    def construct(self, x):
346        input_dtype = self.dtype(x)
347        _check_input_dtype("x", input_dtype, [mstype.float16, mstype.float32], self.cls_name)
348        need_to_reflect = self.less(x, 0.5)
349        neg_input = -x
350        z = self.select(need_to_reflect, neg_input, x - 1)
351
352        @constexpr
353        def _calculate_num_denom(z, k_base_lanczos_coeff, k_lanczos_coefficients):
354            num = 0
355            denom = k_base_lanczos_coeff
356            for i in range(8):
357                j = z + i + 1
358                num = num - k_lanczos_coefficients[i] / (j * j)
359                denom = denom + k_lanczos_coefficients[i] / j
360            return num, denom
361
362        num, denom = _calculate_num_denom(z, self.k_base_lanczos_coeff, self.k_lanczos_coefficients)
363
364        t = z + self.lanczos_gamma_plus_one_half
365        log_t = self.log1p(z / self.lanczos_gamma_plus_one_half) + self.log_lanczos_gamma_plus_one_half
366
367        y = log_t + num / denom - self.k_lanczos_gamma / t
368
369        reduced_input = x + self.abs(self.floor(x + 0.5))
370        reflection = y - self.pi * self.cos(self.pi * reduced_input) / self.sin(self.pi * reduced_input)
371        real_result = self.select(need_to_reflect, reflection, y)
372        nan = F.fill(self.dtype(x), self.shape(x), np.nan)
373
374        return self.select(self.logicaland(self.less(x, 0), self.equal(x, self.floor(x))),
375                           nan, real_result)
376
377
378def _while_helper_func(cond, body, vals):
379    while cond(vals).any():
380        vals = body(vals)
381    return vals
382
383
384def _igamma_series(ax, x, a, enabled):
385    """Helper function for computing Igamma using a power series."""
386
387    logicaland = P.LogicalAnd()
388    greater = P.Greater()
389    shape = P.Shape()
390    dtype = P.DType()
391    select = P.Select()
392
393    # If more data types are supported, this epsilon need to be selected.
394    epsilon = Tensor(np.finfo(np.float32).eps, mstype.float32)
395
396    def cond(vals):
397        enabled = vals[0]
398        return enabled
399
400    def body(vals):
401        enabled = vals[0]
402        r = vals[1]
403        c = vals[2]
404        ans = vals[3]
405        x = vals[4]
406        dc_da = vals[5]
407        dans_da = vals[6]
408
409        r = r + 1
410        dc_da = dc_da * (x / r) + (-1 * c * x) / (r * r)
411        dans_da = dans_da + dc_da
412        c = c * (x / r)
413        ans = ans + c
414        conditional = logicaland(enabled, greater(c / ans, epsilon))
415
416        return (conditional, select(enabled, r, vals[1]),
417                select(enabled, c, vals[2]), select(enabled, ans, vals[3]),
418                select(enabled, x, vals[4]), select(enabled, dc_da, vals[5]),
419                select(enabled, dans_da, vals[6]))
420
421    ones = F.fill(dtype(a), shape(a), 1)
422    zeros = F.fill(dtype(a), shape(a), 0)
423    vals = (enabled, a, ones, ones, x, zeros, zeros)
424
425    vals = _while_helper_func(cond, body, vals)
426    ans = vals[3]
427    return (ans * ax) / a
428
429
430def _igammac_continued_fraction(ax, x, a, enabled):
431    """Helper function for computing Igammac using a continued fraction."""
432
433    abs_x = P.Abs()
434    logicaland = P.LogicalAnd()
435    greater = P.Greater()
436    less = P.Less()
437    notequal = P.NotEqual()
438    shape = P.Shape()
439    dtype = P.DType()
440    select = P.Select()
441
442    # If more data types are supported, this epsilon need to be selected.
443    epsilon = Tensor(np.finfo(np.float32).eps, mstype.float32)
444
445    def cond(vals):
446        enabled = vals[0]
447        c = vals[5]
448        return logicaland(less(c, 2000), enabled)
449
450    def body(vals):
451        enabled = vals[0]
452        ans = vals[1]
453        t = vals[2]
454        y = vals[3]
455        z = vals[4]
456        c = vals[5]
457        pkm1 = vals[6]
458        qkm1 = vals[7]
459        pkm2 = vals[8]
460        qkm2 = vals[9]
461
462        dpkm2_da = vals[10]
463        dqkm2_da = vals[11]
464        dpkm1_da = vals[12]
465        dqkm1_da = vals[13]
466        dans_da = vals[14]
467
468        c = c + 1
469        y = y + 1
470        z = z + 2
471
472        yc = y * c
473        pk = pkm1 * z - pkm2 * yc
474        qk = qkm1 * z - qkm2 * yc
475        qk_is_nonzero = notequal(qk, 0)
476        r = pk / qk
477
478        t = select(qk_is_nonzero, abs_x((ans - r) / r), F.fill(dtype(t), shape(t), 1))
479        ans = select(qk_is_nonzero, r, ans)
480
481        dpk_da = dpkm1_da * z - pkm1 - dpkm2_da * yc + pkm2 * c
482        dqk_da = dqkm1_da * z - qkm1 - dqkm2_da * yc + qkm2 * c
483        dans_da_new = select(qk_is_nonzero, (dpk_da - ans * dqk_da) / qk, dans_da)
484        grad_conditional = select(qk_is_nonzero,
485                                  abs_x(dans_da_new - dans_da),
486                                  F.fill(dtype(dans_da), shape(dans_da), 1))
487
488        pkm2 = pkm1
489        pkm1 = pk
490        qkm2 = qkm1
491        qkm1 = qk
492
493        dpkm2_da = dpkm1_da
494        dqkm2_da = dqkm1_da
495        dpkm1_da = dpk_da
496        dqkm1_da = dqk_da
497
498        rescale = greater(abs_x(pk), 1 / epsilon)
499        pkm2 = select(rescale, pkm2 * epsilon, pkm2)
500        pkm1 = select(rescale, pkm1 * epsilon, pkm1)
501        qkm2 = select(rescale, qkm2 * epsilon, qkm2)
502        qkm1 = select(rescale, qkm1 * epsilon, qkm1)
503
504        dpkm2_da = select(rescale, dpkm2_da * epsilon, dpkm2_da)
505        dqkm2_da = select(rescale, dqkm2_da * epsilon, dqkm2_da)
506        dpkm1_da = select(rescale, dpkm1_da * epsilon, dpkm1_da)
507        dqkm1_da = select(rescale, dqkm1_da * epsilon, dqkm1_da)
508
509        conditional = logicaland(enabled, greater(grad_conditional, epsilon))
510
511        return (conditional, select(enabled, ans, vals[1]), select(enabled, t, vals[2]),
512                select(enabled, y, vals[3]), select(enabled, z, vals[4]),
513                c, select(enabled, pkm1, vals[6]),
514                select(enabled, qkm1, vals[7]), select(enabled, pkm2, vals[8]),
515                select(enabled, qkm2, vals[9]), select(enabled, dpkm2_da, vals[10]),
516                select(enabled, dqkm2_da, vals[11]), select(enabled, dpkm1_da, vals[12]),
517                select(enabled, dqkm1_da, vals[13]), select(enabled, dans_da_new, vals[14]))
518
519    y = 1 - a
520    z = x + y + 1
521    c = F.fill(dtype(x), shape(x), 0)
522    pkm2 = F.fill(dtype(x), shape(x), 1)
523    qkm2 = x
524    pkm1 = x + 1
525    qkm1 = z * x
526    ans = pkm1 / qkm1
527    t = F.fill(dtype(x), shape(x), 1)
528    dpkm2_da = F.fill(dtype(x), shape(x), 0)
529    dqkm2_da = F.fill(dtype(x), shape(x), 0)
530    dpkm1_da = F.fill(dtype(x), shape(x), 0)
531    dqkm1_da = -x
532    dans_da = (dpkm1_da - ans * dqkm1_da) / qkm1
533    vals = (enabled, ans, t, y, z, c, pkm1, qkm1, pkm2, qkm2, dpkm2_da, dqkm2_da, dpkm1_da, dqkm1_da, dans_da)
534    vals = _while_helper_func(cond, body, vals)
535    ans = vals[1]
536    return ans * ax
537
538
539class IGamma(Cell):
540    r"""
541    Calculates lower regularized incomplete Gamma function.
542    The lower regularized incomplete Gamma function is defined as:
543
544    .. math::
545        P(a, x) = gamma(a, x) / Gamma(a) = 1 - Q(a, x)
546
547    where
548
549    .. math::
550        gamma(a, x) = \int_0^x t^{a-1} \exp^{-t} dt
551
552    is the lower incomplete Gamma function.
553
554    Above :math:`Q(a, x)` is the upper regularized complete Gamma function.
555
556    Inputs:
557        - **a** (Tensor) - The input tensor. With float32 data type. `a` should have
558          the same dtype with `x`.
559        - **x** (Tensor) - The input tensor. With float32 data type. `x` should have
560          the same dtype with `a`.
561
562    Outputs:
563        Tensor, has the same dtype as `a` and `x`.
564
565    Raises:
566        TypeError: If dtype of input x and a is not float16 nor float32,
567                   or if x has different dtype with a.
568
569    Supported Platforms:
570        ``Ascend`` ``GPU`` ``CPU``
571
572    Examples:
573        >>> import mindspore as ms
574        >>> import numpy as np
575        >>> a = ms.Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))
576        >>> x = ms.Tensor(np.array([2.0, 3.0, 4.0, 5.0]).astype(np.float32))
577        >>> igamma = ms.nn.IGamma()
578        >>> output = igamma(a, x)
579        >>> print (output)
580        [0.593994  0.35276785  0.21486944  0.13337152]
581    """
582
583    def __init__(self):
584        """Initialize IGamma."""
585        super(IGamma, self).__init__()
586        # const numbers
587        # If more data types are supported, this float max value need to be selected.
588        self.log_maxfloat32 = Tensor(np.log(np.finfo(np.float32).max), mstype.float32)
589
590        # operations
591        self.logicaland = P.LogicalAnd()
592        self.logicalor = P.LogicalOr()
593        self.logicalnot = P.LogicalNot()
594        self.equal = P.Equal()
595        self.greater = P.Greater()
596        self.less = P.Less()
597        self.neg = P.Neg()
598        self.log = P.Log()
599        self.exp = P.Exp()
600        self.select = P.Select()
601        self.zeroslike = P.ZerosLike()
602        self.shape = P.Shape()
603        self.dtype = P.DType()
604        self.lgamma = LGamma()
605        self.cast = P.Cast()
606
607    def construct(self, a, x):
608        a_dtype = self.dtype(a)
609        x_dtype = self.dtype(x)
610        _check_input_dtype("a", a_dtype, [mstype.float32], self.cls_name)
611        _check_input_dtype("x", x_dtype, a_dtype, self.cls_name)
612        domain_error = self.logicalor(self.less(x, 0), self.less(a, 0))
613        use_igammac = self.logicaland(self.greater(x, 1), self.greater(x, a))
614        ax = a * self.log(x) - x - self.lgamma(a)
615        para_shape = self.shape(ax)
616        if para_shape != ():
617            x = F.broadcast_to(x, para_shape)
618            a = F.broadcast_to(a, para_shape)
619        x_is_zero = self.equal(x, 0)
620        underflow = self.less(ax, self.neg(self.log_maxfloat32))
621        ax = self.exp(ax)
622        enabled = self.logicalnot(self.logicalor(self.logicalor(x_is_zero, domain_error), underflow))
623        output = self.select(use_igammac,
624                             1 - _igammac_continued_fraction(ax, x, a, self.logicaland(enabled, use_igammac)),
625                             _igamma_series(ax, x, a, self.logicaland(enabled, self.logicalnot(use_igammac))))
626        output = self.select(x_is_zero, self.zeroslike(output), output)
627        output = self.select(domain_error, F.fill(self.dtype(a), self.shape(a), np.nan), output)
628        return output
629
630
631class LBeta(Cell):
632    r"""
633    This method avoids the numeric cancellation by explicitly
634    decomposing lgamma into the Stirling approximation and an explicit log_gamma_correction, and cancelling
635    the large terms from the Striling analytically.
636
637    This is semantically equal to
638
639    .. math::
640        P(x, y) = lgamma(x) + lgamma(y) - lgamma(x + y).
641
642    The method is more accurate for arguments above 8. The reason for accuracy loss in the naive computation
643    is catastrophic cancellation between the lgammas.
644
645    Inputs:
646        - **x** (Tensor) - The input tensor. With float16 or float32 data type. `x` should have
647          the same dtype with `y`.
648        - **y** (Tensor) - The input tensor. With float16 or float32 data type. `y` should have
649          the same dtype with `x`.
650
651    Outputs:
652        Tensor, has the same dtype as `x` and `y`.
653
654    Raises:
655        TypeError: If dtype of `x` or `y` is neither float16 nor float32,
656                   or if `x` has different dtype with `y`.
657
658    Supported Platforms:
659        ``Ascend`` ``GPU``
660
661    Examples:
662        >>> import mindspore as ms
663        >>> import numpy as np
664        >>> x = ms.Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))
665        >>> y = ms.Tensor(np.array([2.0, 3.0, 14.0, 15.0]).astype(np.float32))
666        >>> lbeta = ms.nn.LBeta()
667        >>> output = lbeta(y, x)
668        >>> print(output)
669        [-1.7917596  -4.094345  -12.000229  -14.754799]
670    """
671
672    def __init__(self):
673        """Initialize LBeta."""
674        super(LBeta, self).__init__()
675        # const numbers
676        self.log_2pi = np.log(2 * np.pi)
677        self.minimax_coeff = [-0.165322962780713e-02,
678                              0.837308034031215e-03,
679                              -0.595202931351870e-03,
680                              0.793650666825390e-03,
681                              -0.277777777760991e-02,
682                              0.833333333333333e-01]
683
684        # operations
685        self.log = P.Log()
686        self.log1p = P.Log1p()
687        self.less = P.Less()
688        self.select = P.Select()
689        self.shape = P.Shape()
690        self.dtype = P.DType()
691        self.lgamma = LGamma()
692        self.const = P.ScalarToTensor()
693
694    def construct(self, x, y):
695        x_dtype = self.dtype(x)
696        y_dtype = self.dtype(y)
697        _check_input_dtype("x", x_dtype, [mstype.float16, mstype.float32], self.cls_name)
698        _check_input_dtype("y", y_dtype, x_dtype, self.cls_name)
699        x_plus_y = x + y
700        para_shape = self.shape(x_plus_y)
701        if para_shape != ():
702            x = F.broadcast_to(x, para_shape)
703            y = F.broadcast_to(y, para_shape)
704        comp_less = self.less(x, y)
705        x_min = self.select(comp_less, x, y)
706        y_max = self.select(comp_less, y, x)
707
708        @constexpr
709        def _log_gamma_correction(x, minimax_coeff):
710            inverse_x = 1. / x
711            inverse_x_squared = inverse_x * inverse_x
712            accum = minimax_coeff[0]
713            for i in range(1, 6):
714                accum = accum * inverse_x_squared + minimax_coeff[i]
715            return accum * inverse_x
716
717        log_gamma_correction_x = _log_gamma_correction(x_min, self.minimax_coeff)
718        log_gamma_correction_y = _log_gamma_correction(y_max, self.minimax_coeff)
719        log_gamma_correction_x_y = _log_gamma_correction(x_plus_y, self.minimax_coeff)
720
721        # Two large arguments case: y >= x >= 8.
722        log_beta_two_large = self.const(0.5 * self.log_2pi, x_dtype) - 0.5 * self.log(y_max) \
723                             + log_gamma_correction_x + log_gamma_correction_y - log_gamma_correction_x_y \
724                             + (x_min - 0.5) * self.log(x_min / (x_min + y_max)) - y_max * self.log1p(x_min / y_max)
725
726        cancelled_stirling = -1 * (x_min + y_max - 0.5) * self.log1p(x_min / y_max) - x_min * self.log(y_max) + x_min
727        correction = log_gamma_correction_y - log_gamma_correction_x_y
728        log_gamma_difference_big_y = correction + cancelled_stirling
729
730        # One large argument case: x < 8, y >= 8.
731        log_beta_one_large = self.lgamma(x_min) + log_gamma_difference_big_y
732
733        # Small arguments case: x <= y < 8.
734        log_beta_small = self.lgamma(x_min) + self.lgamma(y_max) - self.lgamma(x_min + y_max)
735        comp_xless8 = self.less(x_min, 8)
736        comp_yless8 = self.less(y_max, 8)
737        temp = self.select(comp_yless8, log_beta_small, log_beta_one_large)
738        return self.select(comp_xless8, temp, log_beta_two_large)
739
740
741@_primexpr
742def get_broadcast_matmul_shape(x_shape, y_shape, prim_name=None):
743    """get broadcast_matmul shape"""
744    msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
745
746    def _check_len():
747        if (len(x_shape) < 2) or (len(y_shape) < 2):
748            raise ValueError(f"{msg_prefix} length of 'x_shape' and 'y_shape' must be equal to or greater than 2, "
749                             f"but got the length of 'x_shape': {len(x_shape)} and the length of 'y_shape': "
750                             f"{len(y_shape)}.")
751    _check_len()
752    x_shape_batch = x_shape[:-2]
753    y_shape_batch = y_shape[:-2]
754    if x_shape_batch == y_shape_batch:
755        return x_shape, y_shape
756    x_len = len(x_shape)
757    y_len = len(y_shape)
758    length = x_len if x_len < y_len else y_len
759    broadcast_shape_back = []
760
761    def _check_broadcast(x_val, y_val, i):
762        if not (x_val == 1 or y_val == 1 or x_val == y_val):
763            raise ValueError(f"{msg_prefix} 'x_shape[{i}]' must be equal to 1, or the 'y_shape[{i}]' must be equal "
764                             f"to 1, or the 'x_shape[{i}]' must be equal to 'y_shape[{i}]', but got "
765                             f"'x_shape[{i}]': {x_val}, 'y_shape[{i}]': {y_val}.")
766
767    for i in range(-length, -2):
768        _check_broadcast(x_shape[i], y_shape[i], i)
769        if x_shape[i] == 1:
770            broadcast_shape_back.append(y_shape[i])
771        elif y_shape[i] == 1:
772            broadcast_shape_back.append(x_shape[i])
773        else:
774            broadcast_shape_back.append(x_shape[i])
775
776    broadcast_shape_front = y_shape[0: y_len - length] if length == x_len else x_shape[0: x_len - length]
777    x_broadcast_shape = broadcast_shape_front + tuple(broadcast_shape_back) + x_shape[-2:]
778    y_broadcast_shape = broadcast_shape_front + tuple(broadcast_shape_back) + y_shape[-2:]
779    return x_broadcast_shape, y_broadcast_shape
780
781
782@_primexpr
783def check_col_row_equal(x1_shape, x2_shape, transpose_x1, transpose_x2, prim_name=None):
784    """check col and row equal"""
785    def _check(x1_col, x2_row):
786        msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
787        if x1_col != x2_row:
788            raise ValueError(f"{msg_prefix} column of matrix dimensions of 'x1' must be equal to "
789                             f"the row of matrix dimensions of 'x2', but got 'x1_col' {x1_col} and 'x2_row' {x2_row}.")
790
791    if len(x1_shape) == 1:
792        transpose_x1 = False
793        x1_shape = (1,) + x1_shape
794    if len(x2_shape) == 1:
795        transpose_x2 = False
796        x2_shape = x2_shape + (1,)
797    x1_last = x1_shape[-2:]
798    x2_last = x2_shape[-2:]
799    x1_col = x1_last[not transpose_x1]  # x1_col = x1_last[1] if (not transpose_a) else x1_last[0]
800    x2_row = x2_last[transpose_x2]  # x2_row = x2_last[0] if (not transpose_b) else x2_last[1]
801    _check(x1_col, x2_row)
802
803
804def matmul_op_select(x1_shape, x2_shape, transpose_x1, transpose_x2):
805    """select matmul op"""
806    x1_dim, x2_dim = len(x1_shape), len(x2_shape)
807    if x1_dim == 1 and x2_dim == 1:
808        matmul_op = P.Mul()
809    elif x1_dim <= 2 and x2_dim <= 2:
810        transpose_x1 = False if x1_dim == 1 else transpose_x1
811        transpose_x2 = False if x2_dim == 1 else transpose_x2
812        matmul_op = P.MatMul(transpose_x1, transpose_x2)
813    elif x1_dim == 1 and x2_dim > 2:
814        matmul_op = P.BatchMatMul(False, transpose_x2)
815    elif x1_dim > 2 and x2_dim == 1:
816        matmul_op = P.BatchMatMul(transpose_x1, False)
817    else:
818        matmul_op = P.BatchMatMul(transpose_x1, transpose_x2)
819    return matmul_op
820
821
822class MatMul(Cell):
823    r"""
824    The nn.MatMul interface is deprecated, please use the :class:`mindspore.ops.matmul` instead.
825
826    Supported Platforms:
827        deprecated
828    """
829
830    @deprecated('1.2', 'ops.matmul', False)
831    def __init__(self, transpose_x1=False, transpose_x2=False):
832        """Initialize MatMul."""
833        super(MatMul, self).__init__()
834
835        validator.check_value_type('transpose_x1', transpose_x1, [bool], self.cls_name)
836        validator.check_value_type('transpose_x2', transpose_x2, [bool], self.cls_name)
837        self.transpose_x1 = transpose_x1
838        self.transpose_x2 = transpose_x2
839        self.shape_op = P.Shape()
840        self.expand_op = P.ExpandDims()
841        self.squeeze_left_op = P.Squeeze(-2)
842        self.squeeze_right_op = P.Squeeze(-1)
843        self.reduce_sum_op = P.ReduceSum(keep_dims=False)
844
845    def construct(self, x1, x2):
846        x1_shape = self.shape_op(x1)
847        x2_shape = self.shape_op(x2)
848        check_col_row_equal(x1_shape, x2_shape, self.transpose_x1, self.transpose_x2, self.cls_name)
849        matmul_op = matmul_op_select(x1_shape, x2_shape, self.transpose_x1, self.transpose_x2)
850
851        x1_dim, x2_dim = len(x1_shape), len(x2_shape)
852        if x1_dim == x2_dim and x2_dim == 1:
853            return self.reduce_sum_op(matmul_op(x1, x2), -1)
854        if x1_dim == 1:
855            x1 = self.expand_op(x1, 0)
856            x1_shape = self.shape_op(x1)
857        if x2_dim == 1:
858            x2 = self.expand_op(x2, 1)
859            x2_shape = self.shape_op(x2)
860
861        x1_broadcast_shape, x2_broadcast_shape = get_broadcast_matmul_shape(x1_shape, x2_shape)
862        if x1_broadcast_shape != x1_shape:
863            x1 = F.broadcast_to(x1, x1_broadcast_shape)
864        if x2_broadcast_shape != x2_shape:
865            x2 = F.broadcast_to(x2, x2_broadcast_shape)
866
867        matmul_broadcast = matmul_op(x1, x2)
868
869        if x1_dim == 1:
870            matmul_broadcast = self.squeeze_left_op(matmul_broadcast)
871        if x2_dim == 1:
872            matmul_broadcast = self.squeeze_right_op(matmul_broadcast)
873
874        return matmul_broadcast
875
876
877class CosineSimilarity(Cell):
878    r"""
879    Computes cosine similarity.
880
881    .. math::
882        \mathcal{K} = \frac{\textbf{x}\textbf{y}^{\top}}{\parallel \textbf{x} \parallel \parallel \textbf{y} \parallel},
883
884    where :math:`\mathcal{K}` is the similarity, :math:`\textbf{x}` is the first tensor `x1`,
885    :math:`\textbf{y}` is the second tensor `x2`.
886
887    To avoid numerical errors when dividing by small numbers,
888    the lower bound of :math:`\parallel \textbf{x} \parallel \parallel \textbf{y} \parallel` is set to `eps`.
889
890    Args:
891        dim (int, optional): Dimension. Default: ``1`` .
892        eps (float, optional): Small value. Default: ``1e-08`` .
893
894    Inputs:
895        - **x1** (Tensor) - The first tensor :math:`\textbf{x}`.
896          Shape: :math:`(\ast_1, D, \ast_2)` where :math:`D` is at position `dim`.
897        - **x2** (Tensor) - The second tensor :math:`\textbf{y}`. The shape is the same as `x1`.
898
899    Outputs:
900        Tensor, with shape :math:`(\ast_1, \ast_2)`, the data type will be inferred automatically.
901
902    Raises:
903        TypeError: If `x1` or `x2` is not a Tensor.
904
905    Supported Platforms:
906        ``Ascend`` ``GPU`` ``CPU``
907
908    Examples:
909        >>> import mindspore as ms
910        >>> import numpy as np
911        >>> x1 = ms.Tensor([[1.0, 3.0, 4.0, 7.0], [2.0, 4.0, 2.0, 5.0], [3.0, 1.0, 5.0, 8.0]])
912        >>> x2 = ms.Tensor([[2.0, 4.0, 2.0, 5.0], [3.0, 1.0, 5.0, 8.0], [1.0, 3.0, 4.0, 7.0]])
913        >>> func = ms.nn.layer.CosineSimilarity()
914        >>> out = func(x1, x2)
915        >>> print(out.asnumpy())
916        [0.9402562 0.8614609 0.9516245]
917    """
918
919    def __init__(self, dim=1, eps=1e-08):
920        """Initialize CosineSimilarity."""
921        super().__init__()
922        self.dim = dim
923        self.eps = eps
924        self.mul = P.Mul()
925        self.div = P.DivNoNan()
926        self.maximum = P.Maximum()
927        self.cast = P.Cast()
928
929    def construct(self, x1, x2):
930        if not isinstance(x1, Tensor):
931            raise TypeError(f"For 'CosineSimilarity', 'x1' must be a tensor, but got {type(x1)}")
932        if not isinstance(x2, Tensor):
933            raise TypeError(f"For 'CosineSimilarity', 'x2' must be a tensor, but got {type(x2)}")
934        w12 = self.mul(x1, x2).sum(self.dim)
935        w1 = self.mul(x1, x1).sum(self.dim)
936        w2 = self.mul(x2, x2).sum(self.dim)
937        n12 = self.maximum(self.mul(w1, w2), self.eps * self.eps).sqrt()
938        out = self.div(w12, n12)
939        return out
940
941
942class Moments(Cell):
943    """
944    The Moments class will be deprecated in the future,
945    this function can be replaced by :func:`ops.var_mean`
946    """
947
948    def __init__(self, axis=None, keep_dims=None):
949        """Initialize Moments."""
950        super(Moments, self).__init__()
951        logger.warning("The Moments class will be deprecated in the future,"
952                       "this function can be replaced by :func:`ops.var_mean`")
953        if axis is None:
954            axis = ()
955        if isinstance(axis, tuple):
956            for idx, item in enumerate(axis):
957                validator.check_value_type("axis[%d]" % idx, item, [int], self.cls_name)
958        self.axis = validator.check_value_type('axis', axis, [int, tuple], self.cls_name)
959        if keep_dims is None:
960            keep_dims = False
961        self.keep_dims = validator.check_value_type('keep_dims', keep_dims, [bool], self.cls_name)
962        self.cast = P.Cast()
963        self.reduce_mean = P.ReduceMean(keep_dims=True)
964        self.square_diff = P.SquaredDifference()
965        self.squeeze = P.Squeeze(self.axis)
966
967    def construct(self, x):
968        tensor_dtype = F.dtype(x)
969        _check_input_dtype("input x", tensor_dtype, [mstype.float16, mstype.float32], self.cls_name)
970        if tensor_dtype == mstype.float16:
971            x = self.cast(x, mstype.float32)
972        mean = self.reduce_mean(x, self.axis)
973        variance = self.reduce_mean(self.square_diff(x, F.stop_gradient(mean)), self.axis)
974        if not self.keep_dims:
975            mean = self.squeeze(mean)
976            variance = self.squeeze(variance)
977        if tensor_dtype == mstype.float16:
978            mean = self.cast(mean, mstype.float16)
979            variance = self.cast(variance, mstype.float16)
980            return mean, variance
981        return mean, variance
982
983
984class MatInverse(Cell):
985    """
986    Calculates the inverse of Positive-Definite Hermitian matrix using Cholesky decomposition.
987
988    Inputs:
989        - **x** (Tensor[Number]) - The input tensor. It must be a positive-definite matrix.
990          With float16 or float32 data type.
991
992    Outputs:
993        Tensor, has the same dtype as the `x`.
994
995    Raises:
996        TypeError: If dtype of `x` is neither float16 nor float32.
997
998    Supported Platforms:
999        ``GPU``
1000
1001    Examples:
1002        >>> import mindspore as ms
1003        >>> import numpy as np
1004        >>> x = ms.Tensor(np.array([[4, 12, -16], [12, 37, -43], [-16, -43, 98]]).astype(np.float32))
1005        >>> op = ms.nn.MatInverse()
1006        >>> output = op(x)
1007        >>> print(output)
1008        [[49.36112  -13.555558  2.1111116]
1009         [-13.555558  3.7777784  -0.5555557]
1010         [2.1111116  -0.5555557  0.11111113]]
1011    """
1012
1013    def __init__(self):
1014        """Initialize MatInverse."""
1015        super(MatInverse, self).__init__()
1016        self.dtype = P.DType()
1017        self.choleskytrsm = P.CholeskyTrsm()
1018        self.matmul = MatMul(transpose_x1=True)
1019
1020    def construct(self, a):
1021        input_dtype = self.dtype(a)
1022        _check_input_dtype("input_a", input_dtype, [mstype.float16, mstype.float32], self.cls_name)
1023        l_inverse = self.choleskytrsm(a)
1024        a_inverse = self.matmul(l_inverse, l_inverse)
1025        return a_inverse
1026
1027
1028class MatDet(Cell):
1029    """
1030    Calculates the determinant of Positive-Definite Hermitian matrix using Cholesky decomposition.
1031
1032    Inputs:
1033        - **x** (Tensor[Number]) - The input tensor. It must be a positive-definite matrix.
1034          With float16 or float32 data type.
1035
1036    Outputs:
1037        Tensor, has the same dtype as the `x`.
1038
1039    Raises:
1040        TypeError: If dtype of `x` is neither float16 nor float32.
1041
1042    Supported Platforms:
1043        ``GPU``
1044
1045    Examples:
1046        >>> import mindspore as ms
1047        >>> import numpy as np
1048        >>> x = ms.Tensor(np.array([[4, 12, -16], [12, 37, -43], [-16, -43, 98]]).astype(np.float32))
1049        >>> op = ms.nn.MatDet()
1050        >>> output = op(x)
1051        >>> print(output)
1052        35.999996
1053    """
1054
1055    def __init__(self):
1056        """Initialize MatDet."""
1057        super(MatDet, self).__init__()
1058        self.dtype = P.DType()
1059        self.cholesky = P.Cholesky()
1060        self.det_triangle = P.DetTriangle()
1061        self.square = P.Square()
1062
1063    def construct(self, a):
1064        input_dtype = self.dtype(a)
1065        _check_input_dtype("input_a", input_dtype, [mstype.float16, mstype.float32], self.cls_name)
1066        l = self.cholesky(a)
1067        l_det = self.det_triangle(l)
1068        a_det = self.square(l_det)
1069        return a_det
1070