From 1f06e44743e845ad491f7f253022e9fe6df2e53f Mon Sep 17 00:00:00 2001 From: ddunleavy Date: Tue, 18 Jun 2024 10:45:05 -0700 Subject: [PATCH] Move `tensorflow/tsl/protobuf` to `xla/tsl/protobuf` PiperOrigin-RevId: 644440378 --- .../model_server_clients/tensorflow_serving_client.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tfx/components/infra_validator/model_server_clients/tensorflow_serving_client.py b/tfx/components/infra_validator/model_server_clients/tensorflow_serving_client.py index 52fc79b8a81..c5a365fe7cf 100644 --- a/tfx/components/infra_validator/model_server_clients/tensorflow_serving_client.py +++ b/tfx/components/infra_validator/model_server_clients/tensorflow_serving_client.py @@ -37,7 +37,8 @@ 'the TF version doesn\'t match with the TF Serving version. ' 'We will try importing again with a workaround:%s', err) from tensorflow.core.protobuf import error_codes_pb2 as old_error_codes_pb2 - from tensorflow.tsl.protobuf import error_codes_pb2 as new_error_codes_pb2 + from tensorflow.compiler.xla.tsl.protobuf import error_codes_pb2 as new_error_codes_pb2 + old_error_codes_pb2._CODE = new_error_codes_pb2._CODE # pylint: disable=protected-access # pytype: disable=module-attr # Retry.