Change get_replicated_var_handle to accept resource tensors instead of variables

This avoid the needs to wrap packed tensor handle in a variable. This enables a
ongoing refactoring of DistributedVariable.

PiperOrigin-RevId: 371530563

Author: crccw

mesh_tensorflow/tpu_variables.py

@@ -21,12 +21,14 @@
 
 import contextlib
 
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.compat import compat
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import gen_resource_variable_ops
 
 try:
-  from tensorflow.python.types import core  # pylint:disable=g-import-not-at-top,g-direct-tensorflow-import
+  from tensorflow.python.types import core  # pylint:disable=g-import-not-at-top
   TF_23 = True
 except ImportError:
   TF_23 = False
@@ -80,8 +82,11 @@ def handle(self):
     tpu_context = _enclosing_tpu_context()
     if tpu_context is None:
       return self._primary_var.handle
-
-    return tpu_context.get_replicated_var_handle(self._name, self._vars)
+    if compat.forward_compatible(2021, 4, 29):
+      handles = [v.handle for v in self._vars]
+      return tpu_context.get_replicated_var_handle(self._name, handles)
+    else:
+      return tpu_context.get_replicated_var_handle(self._name, self._vars)
 
   @contextlib.contextmanager
   def _assign_dependencies(self):