@@ -29,7 +29,6 @@ limitations under the License.
29
29
#include " tensorflow/core/framework/tensor.h"
30
30
#include " tensorflow/core/framework/tensor_shape.h"
31
31
#include " tensorflow/core/framework/tensor_types.h"
32
- #include " tensorflow/core/platform/logging.h"
33
32
#include " tensorflow/core/platform/status.h"
34
33
#include " tensorflow/core/platform/types.h"
35
34
#include " tensorflow_compression/cc/lib/bit_coder.h"
@@ -38,7 +37,6 @@ namespace tensorflow_compression {
38
37
namespace {
39
38
namespace errors = tensorflow::errors;
40
39
using tensorflow::DEVICE_CPU;
41
- using tensorflow::FromAbslStatus;
42
40
using tensorflow::OpKernel;
43
41
using tensorflow::OpKernelConstruction;
44
42
using tensorflow::OpKernelContext;
@@ -48,6 +46,13 @@ using tensorflow::TensorShape;
48
46
using tensorflow::TensorShapeUtils;
49
47
using tensorflow::tstring;
50
48
49
+ #define OP_REQUIRES_OK_ABSL (context, status ) \
50
+ { \
51
+ auto s = (status); \
52
+ OP_REQUIRES (context, s.ok (), tensorflow::Status ( \
53
+ static_cast <tensorflow::error::Code>(s.code ()), s.message ())); \
54
+ }
55
+
51
56
class RunLengthGammaEncodeOp : public OpKernel {
52
57
public:
53
58
explicit RunLengthGammaEncodeOp (OpKernelConstruction* context)
@@ -142,7 +147,7 @@ class RunLengthGammaDecodeOp : public OpKernel {
142
147
for (int64_t i = 0 ; i < data.size (); i++) {
143
148
// Get number of zeros.
144
149
auto num_zeros = dec.ReadGamma ();
145
- OP_REQUIRES (context, num_zeros.ok (), FromAbslStatus (num_zeros. status () ));
150
+ OP_REQUIRES_OK_ABSL (context, num_zeros.status ());
146
151
147
152
// Advance the index to the next non-zero element.
148
153
i += *num_zeros - 1 ;
@@ -157,11 +162,11 @@ class RunLengthGammaDecodeOp : public OpKernel {
157
162
158
163
// Get sign of value.
159
164
auto positive = dec.ReadOneBit ();
160
- OP_REQUIRES (context, positive.ok (), FromAbslStatus (positive. status () ));
165
+ OP_REQUIRES_OK_ABSL (context, positive.status ());
161
166
162
167
// Get magnitude.
163
168
auto magnitude = dec.ReadGamma ();
164
- OP_REQUIRES (context, magnitude.ok (), FromAbslStatus (magnitude. status () ));
169
+ OP_REQUIRES_OK_ABSL (context, magnitude.status ());
165
170
166
171
// Write value to data tensor element at index.
167
172
data (i) = *positive ? *magnitude : -*magnitude;
@@ -172,5 +177,7 @@ class RunLengthGammaDecodeOp : public OpKernel {
172
177
REGISTER_KERNEL_BUILDER (Name(" RunLengthGammaDecode" ).Device(DEVICE_CPU),
173
178
RunLengthGammaDecodeOp);
174
179
180
+ #undef OP_REQUIRES_OK_ABSL
181
+
175
182
} // namespace
176
183
} // namespace tensorflow_compression
0 commit comments