Skip to content

Commit 50ad094

Browse files
Johannes Ballécopybara-github
authored andcommitted
Simplified error status conversion for backwards compatibility with TF 2.10.
PiperOrigin-RevId: 472724609 Change-Id: I9b6be102add6fc82cc8b7c7fcf2f2cb16219e1ca
1 parent 6d3638c commit 50ad094

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

tensorflow_compression/cc/kernels/run_length_gamma_kernels.cc

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ limitations under the License.
2929
#include "tensorflow/core/framework/tensor.h"
3030
#include "tensorflow/core/framework/tensor_shape.h"
3131
#include "tensorflow/core/framework/tensor_types.h"
32-
#include "tensorflow/core/platform/logging.h"
3332
#include "tensorflow/core/platform/status.h"
3433
#include "tensorflow/core/platform/types.h"
3534
#include "tensorflow_compression/cc/lib/bit_coder.h"
@@ -38,7 +37,6 @@ namespace tensorflow_compression {
3837
namespace {
3938
namespace errors = tensorflow::errors;
4039
using tensorflow::DEVICE_CPU;
41-
using tensorflow::FromAbslStatus;
4240
using tensorflow::OpKernel;
4341
using tensorflow::OpKernelConstruction;
4442
using tensorflow::OpKernelContext;
@@ -48,6 +46,13 @@ using tensorflow::TensorShape;
4846
using tensorflow::TensorShapeUtils;
4947
using tensorflow::tstring;
5048

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+
5156
class RunLengthGammaEncodeOp : public OpKernel {
5257
public:
5358
explicit RunLengthGammaEncodeOp(OpKernelConstruction* context)
@@ -142,7 +147,7 @@ class RunLengthGammaDecodeOp : public OpKernel {
142147
for (int64_t i = 0; i < data.size(); i++) {
143148
// Get number of zeros.
144149
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());
146151

147152
// Advance the index to the next non-zero element.
148153
i += *num_zeros - 1;
@@ -157,11 +162,11 @@ class RunLengthGammaDecodeOp : public OpKernel {
157162

158163
// Get sign of value.
159164
auto positive = dec.ReadOneBit();
160-
OP_REQUIRES(context, positive.ok(), FromAbslStatus(positive.status()));
165+
OP_REQUIRES_OK_ABSL(context, positive.status());
161166

162167
// Get magnitude.
163168
auto magnitude = dec.ReadGamma();
164-
OP_REQUIRES(context, magnitude.ok(), FromAbslStatus(magnitude.status()));
169+
OP_REQUIRES_OK_ABSL(context, magnitude.status());
165170

166171
// Write value to data tensor element at index.
167172
data(i) = *positive ? *magnitude : -*magnitude;
@@ -172,5 +177,7 @@ class RunLengthGammaDecodeOp : public OpKernel {
172177
REGISTER_KERNEL_BUILDER(Name("RunLengthGammaDecode").Device(DEVICE_CPU),
173178
RunLengthGammaDecodeOp);
174179

180+
#undef OP_REQUIRES_OK_ABSL
181+
175182
} // namespace
176183
} // namespace tensorflow_compression

0 commit comments

Comments
 (0)