From d32bfbeedcd6c657af91b9607d19cd0cc4e481ad Mon Sep 17 00:00:00 2001 From: Brian Wieder Date: Fri, 10 Mar 2023 11:34:29 -0800 Subject: [PATCH] Move `convert_to_tensor`, `convert_to_tensor_v1`, `convert_to_tensor_v1_with_dispatch`, `convert_to_tensor_v2_with_dispatch`, and `convert_to_tensor_v2` into `tensor_conversion_registry`. Also renaming `tensor_conversion_registry` to `tensor_conversion` to match. PiperOrigin-RevId: 515693810 --- mesh_tensorflow/tpu_variables.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mesh_tensorflow/tpu_variables.py b/mesh_tensorflow/tpu_variables.py index e5b1ba2d..01ee6d3f 100644 --- a/mesh_tensorflow/tpu_variables.py +++ b/mesh_tensorflow/tpu_variables.py @@ -24,7 +24,7 @@ # pylint: disable=g-direct-tensorflow-import from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_conversion_registry +from tensorflow.python.framework import tensor_conversion from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gen_resource_variable_ops @@ -224,7 +224,7 @@ def _tensor_conversion(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access -tensor_conversion_registry.register_tensor_conversion_function( +tensor_conversion.register_tensor_conversion_function( ReplicatedVariable, _tensor_conversion) if not TF_23: