-
Notifications
You must be signed in to change notification settings - Fork 74k
/
cwise_ops_gradients.h
210 lines (184 loc) · 6.86 KB
/
cwise_ops_gradients.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_CWISE_OPS_GRADIENTS_H_
#define TENSORFLOW_CORE_KERNELS_CWISE_OPS_GRADIENTS_H_
#define EIGEN_USE_THREADS
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/kernels/cwise_ops.h"
namespace Eigen {
namespace internal {
// Gradient for the tanh function
template <typename T>
struct scalar_tanh_gradient_op {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const T
operator()(const T& output, const T& output_gradient) const {
return output_gradient * (T(1) - output * output);
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet
packetOp(const Packet& output, const Packet& output_gradient) const {
return pmul(output_gradient,
psub(pset1<Packet>(T(1)), pmul(output, output)));
}
};
template <typename T>
struct functor_traits<scalar_tanh_gradient_op<T>> {
enum {
Cost = NumTraits<T>::AddCost + 2 * NumTraits<T>::MulCost,
PacketAccess = packet_traits<T>::HasSub && packet_traits<T>::HasMul,
};
};
// Gradient for the sigmoid function
template <typename T>
struct scalar_sigmoid_gradient_op {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const T
operator()(const T& output, const T& output_gradient) const {
return output_gradient * output * (T(1) - output);
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet
packetOp(const Packet& output, const Packet& output_gradient) const {
return pmul(output_gradient,
pmul(output, psub(pset1<Packet>(T(1)), output)));
}
};
template <typename T>
struct functor_traits<scalar_sigmoid_gradient_op<T>> {
enum {
Cost = NumTraits<T>::AddCost + 2 * NumTraits<T>::MulCost,
PacketAccess = packet_traits<T>::HasSub && packet_traits<T>::HasMul,
};
};
// Gradient for the inverse function
template <typename T>
struct scalar_inverse_gradient_op {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const T
operator()(const T& output, const T& output_gradient) const {
if (output_gradient == T(0)) {
return T(0);
} else {
const T out_conj = numext::conj(output);
return -out_conj * out_conj * output_gradient;
}
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet
packetOp(const Packet& output, const Packet& output_gradient) const {
const Packet out_conj = pconj(output);
return mul_no_nan_op<T>().packetOp(pnegate(pmul(out_conj, out_conj)),
output_gradient);
}
};
template <typename T>
struct functor_traits<scalar_inverse_gradient_op<T>> {
enum {
Cost = NumTraits<T>::AddCost + 2 * NumTraits<T>::MulCost,
PacketAccess = packet_traits<T>::HasMul,
};
};
// Gradient for the sqrt function
template <typename T>
struct scalar_sqrt_gradient_op {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const T
operator()(const T& output, const T& output_gradient) const {
if (output_gradient == T(0)) {
return T(0);
} else {
const T out_conj = numext::conj(output);
return (static_cast<T>(0.5) * output_gradient) / out_conj;
}
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet
packetOp(const Packet& output, const Packet& output_gradient) const {
const Packet const_half = pset1<Packet>(static_cast<T>(0.5));
const Packet out_conj = pconj(output);
return mul_no_nan_op<T>().packetOp(pdiv(const_half, out_conj),
output_gradient);
}
};
template <typename T>
struct functor_traits<scalar_sqrt_gradient_op<T>> {
enum {
PacketAccess = packet_traits<T>::HasMul & packet_traits<T>::HasDiv,
Cost = NumTraits<T>::MulCost + scalar_div_cost<T, PacketAccess>::value,
};
};
// Gradient for the rsqrt function
template <typename T>
struct scalar_rsqrt_gradient_op {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const T
operator()(const T& output, const T& output_gradient) const {
if (output_gradient == T(0)) {
return T(0);
} else {
const T out_conj = numext::conj(output);
return static_cast<T>(-0.5) * (output_gradient * out_conj) *
(out_conj * out_conj);
}
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet
packetOp(const Packet& output, const Packet& output_gradient) const {
const Packet const_half = pset1<Packet>(static_cast<T>(-0.5));
const Packet out_conj = pconj(output);
auto safe_pmul = [](const Packet& a, const Packet& b) {
return mul_no_nan_op<T>().packetOp(a, b);
};
return safe_pmul(pmul(const_half, pmul(out_conj, out_conj)),
safe_pmul(out_conj, output_gradient));
}
};
template <typename T>
struct functor_traits<scalar_rsqrt_gradient_op<T>> {
enum {
Cost = 4 * NumTraits<T>::MulCost,
PacketAccess = packet_traits<T>::HasMul,
};
};
} // end namespace internal
} // end namespace Eigen
namespace tensorflow {
namespace functor {
template <typename Device, typename Functor>
struct SimpleBinaryFunctor {
void operator()(const Device& d, typename Functor::tout_type out,
typename Functor::tin_type in0,
typename Functor::tin_type in1);
};
// Partial specialization of BinaryFunctor for CPU devices
typedef Eigen::ThreadPoolDevice CPUDevice;
template <typename Functor>
struct SimpleBinaryFunctor<CPUDevice, Functor> {
void operator()(const CPUDevice& d, typename Functor::tout_type out,
typename Functor::tin_type in0,
typename Functor::tin_type in1) {
out.device(d) = in0.binaryExpr(in1, typename Functor::func());
}
};
template <typename T>
struct tanh_grad : base<T, Eigen::internal::scalar_tanh_gradient_op<T>> {};
template <typename T>
struct sigmoid_grad : base<T, Eigen::internal::scalar_sigmoid_gradient_op<T>> {
};
template <typename T>
struct inverse_grad : base<T, Eigen::internal::scalar_inverse_gradient_op<T>> {
};
template <typename T>
struct sqrt_grad : base<T, Eigen::internal::scalar_sqrt_gradient_op<T>> {};
template <typename T>
struct rsqrt_grad : base<T, Eigen::internal::scalar_rsqrt_gradient_op<T>> {};
template <typename T>
struct igamma_grad_a : base<T, Eigen::internal::scalar_igamma_der_a_op<T>> {};
} // end namespace functor
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_CWISE_OPS_GRADIENTS_H_