1 #pragma once
2
3 #include <c10/util/complex.h>
4 #include <ATen/NumericUtils.h>
5
6 namespace at::native {
7 inline namespace CPU_CAPABILITY {
8
9 // custom min and max to be used in logcumsumexp for complex arguments
10 template <typename scalar_t>
_logcumsumexp_minmax(c10::complex<scalar_t> x,c10::complex<scalar_t> y)11 std::pair<c10::complex<scalar_t>, c10::complex<scalar_t>> _logcumsumexp_minmax(c10::complex<scalar_t> x, c10::complex<scalar_t> y) {
12 if (at::_isnan(y)) { // either real is nan or imag is nan
13 return std::make_pair(y, y);
14 } else if (at::_isnan(x)) { // either real is nan or imag is nan
15 return std::make_pair(x, x);
16 } else {
17 return (x.real() < y.real()) ? std::make_pair(x, y) : std::make_pair(y, x);
18 }
19 }
20
21 template <typename scalar_t>
_log_add_exp_helper(scalar_t x,scalar_t y)22 scalar_t _log_add_exp_helper(scalar_t x, scalar_t y) {
23 // Reference : https://www.tensorflow.org/api_docs/python/tf/math/cumulative_logsumexp
24 scalar_t min = at::_isnan(y) ? y : std::min(x, y); // std::min returns first arg if one of the args is nan
25 scalar_t max = at::_isnan(y) ? y : std::max(x, y); // std::max returns first arg if one of the args is nan
26 if (min != max || std::isfinite(min)) {
27 // nan will be propagated here
28 return std::log1p(std::exp(min - max)) + max;
29 } else {
30 // special case to correctly handle infinite cases
31 return x;
32 }
33 }
34
35 template <typename scalar_t>
_log_add_exp_helper(const c10::complex<scalar_t> & x,const c10::complex<scalar_t> & y)36 c10::complex<scalar_t> _log_add_exp_helper(const c10::complex<scalar_t>& x, const c10::complex<scalar_t>& y) {
37 auto [min, max] = _logcumsumexp_minmax<scalar_t>(x, y);
38 auto min_real = std::real(min);
39 auto max_real = std::real(max);
40
41 if (at::_isnan(min)) { // either real is nan or imag is nan
42 // handling the "infectious" NaNs
43 return {std::numeric_limits<scalar_t>::quiet_NaN(), std::numeric_limits<scalar_t>::quiet_NaN()};
44 } else if (!std::isfinite(min_real) && (min_real == max_real)) {
45 if (min_real < 0) {
46 // handle the -inf case, the imaginary part here does not really matter as the exp(value)
47 // will be around 0.0 and the angle (i.e. the imaginary part) cannot be determined.
48 // It does not matter if we're taking the exp of this value
49 return min;
50 } else {
51 // handle the +inf case, we don't need the special precision for log1p for small values
52 // and to avoid producing nan in case of real(max) == real(min) == +inf
53 return std::log(std::exp(min) + std::exp(max));
54 }
55 } else {
56 return std::log1p(std::exp(min - max)) + max;
57 }
58 }
59
60 } // end namespace
61 } //end at::native
62