Skip to content

Commit

Permalink
Move convert_to_tensor, convert_to_tensor_v1, `convert_to_tensor_…
Browse files Browse the repository at this point in the history
…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
  • Loading branch information
BrianWieder authored and Mesh TensorFlow Team committed Mar 13, 2023
1 parent 4513b6c commit 4861b11
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions mesh_tensorflow/tpu_variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 4861b11

Please sign in to comment.