class_docstr: Class for defining how to save/restore DTensor variable.
class_docstr_tok: ["Class","for","defining","how","to","save\/restore","DTensor","variable","."]
docstr: Restores the same value into all variables.
docstr_tok: ["Restores","the","same","value","into","all","variables","."]
code: def restore(self, restored_tensors, restored_shapes):
      """Restores the same value into all variables."""
      (tensor,) = restored_tensors
  
      @def_function.function
      def _restore(t):
          with ops.device(self._dvariable.device):
              return api.copy_to_mesh(t, self._original_layout)
  
      # This assign establishes connections from restored tensor and tensors
      # being restored to -- so that restore in SPMD can backtrack the DVariable
      # and its layout, given that we're using tf.function style restore.
      # Note that the restored dvaraible is on CPU no matter what as the restoreV2
      # op must run on CPU.
      # TODO(b/159035705): Allow restore for Tensor objects as well?
      # Restore the dvariable back to original layout.
      if self._original_layout.mesh.device_type().upper() != "CPU":
          tensor = _restore(tensor)
      return self._dvariable.assign(
          math_ops.cast(tensor, dtype=self._dvariable.dtype)
          if self._dvariable.save_as_bf16
          else tensor
      )
  
imports: import functools
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import layout
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.trackable import base
  from tensorflow.python.training.saving import saveable_object
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 3 for calling the _DVariableSaveable obj's underlying member method code with example usage: obj.restore(restored_tensors, restored_shapes) and returns: self, api