Skip to content
This repository was archived by the owner on Jan 21, 2025. It is now read-only.

Commit 4513b6c

Browse files
jam14jMesh TensorFlow Team
authored and
Mesh TensorFlow Team
committed
Remove legacy references from ops.py.
This is done to eventually remove the lazy loads in `indexed_slices.py`. PiperOrigin-RevId: 504679098
1 parent a9a13a8 commit 4513b6c

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

mesh_tensorflow/tpu_variables.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424

2525
# pylint: disable=g-direct-tensorflow-import
2626
from tensorflow.python.framework import ops
27+
from tensorflow.python.framework import tensor_conversion_registry
2728
from tensorflow.python.ops import control_flow_ops
2829
from tensorflow.python.ops import gen_resource_variable_ops
2930

@@ -223,7 +224,8 @@ def _tensor_conversion(var, dtype=None, name=None, as_ref=False):
223224
return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access
224225

225226

226-
ops.register_tensor_conversion_function(ReplicatedVariable, _tensor_conversion)
227+
tensor_conversion_registry.register_tensor_conversion_function(
228+
ReplicatedVariable, _tensor_conversion)
227229

228230
if not TF_23:
229231
ops.register_dense_tensor_like_type(ReplicatedVariable)

0 commit comments

Comments
 (0)