class_docstr: A replacement for tf.Variable which follows initial value placement.
  
  The class also handles restore/save operations in DTensor. Note that,
  DVariable may fall back to normal tf.Variable at this moment if
  `initial_value` is not a DTensor.
class_docstr_tok: ["A","replacement","for","tf.Variable","which","follows","initial","value","placement",".","The","class","also","handles","restore\/save","operations","in","DTensor",".","Note","that",",","DVariable","may","fall","back","to","normal","tf.Variable","at","this","moment","if","`","initial_value","`","is","not","a","DTensor","."]
docstr: Overrides tf.Variable to fix VarHandleOp placements.
docstr_tok: ["Overrides","tf.Variable","to","fix","VarHandleOp","placements","."]
code: def __init__(self, initial_value, *args, dtype=None, **kwargs):
      """Overrides tf.Variable to fix VarHandleOp placements."""
      # Variables by default use the current device scope for placement. This
      # wrapper has them follow the initial value's placement instead (which will
      # be the DTensor device if the initial value has a layout).
  
      # Pop layout from kwargs since keras make_variable may pass a 'layout'
      # keyword argument. We need to pop it because we are passing kwargs to
      # super class constructor.
      layout = kwargs.pop("layout", None)
      shape = kwargs.get("shape", None)
  
      if callable(initial_value):
          unwrapped = initial_value
          if issubclass(type(initial_value), functools.partial):
              unwrapped = initial_value.func
  
          # If wrapped is a CheckpointInitialValueCallable, this means that
          # we are creating a Variable during a checkpoint restore.
          # Thus the restore will happen now through this callable
          # and we will create the DVariable with the restored dtensor.
          if issubclass(
              type(unwrapped), trackable.CheckpointInitialValueCallable
          ):
              if not shape or not layout:
                  raise ValueError(
                      "Expected shape and layout to be not None."
                  )
  
              # CheckpointInitialValueCallable will call an eager tf.RestoreV2,
              # which does not have any shape information or layout information
              # attached. Thus we will do two things to have them correctly specified:
              #
              # The default layout scope allows us to correctly specify the output
              # layout of the tf.RestoreV2 that will be called
              #
              # Passing shard_info with the correct shape allows the tf.RestoreV2
              # ShapeInference to extract the shape.
              initial_value = api.call_with_layout(
                  initial_value,
                  layout,
                  shard_info=trackable.ShardInfo(
                      shape=shape, offset=[0] * len(shape)
                  ),
              )
          else:
              initial_value = initial_value()
  
      # When the initial value came from a Checkpoint restoration, fetch tensor.
      if isinstance(initial_value, trackable.CheckpointInitialValue):
          initial_value = initial_value.wrapped_value
  
      initial_value = ops.convert_to_tensor(initial_value, dtype=dtype)
      variable_device = initial_value.device
      self._save_as_bf16 = False
      # TODO(b/159035705): The following code enables variable creation inside
      # a tf.function. However, it requires a global dtensor device.
      # if not variable_device and not tf.executing_eagerly():
      #   try:
      #     initial_value.op.get_attr("_layout")
      #   except ValueError:
      #     pass
      #   else:
      #     # The initial value is a DTensor, but because the DTensor device is
      #     # only active during eager execution at the moment we need to
      #     # translate that into a placement for the eager VarHandleOp.
      #     variable_device = _dtensor_device().name
      with ops.device(variable_device):
          # If initial tensor assigned to DVariable is DTensor, record the layout of
          # the resource so that this can be queried.
          if context.executing_eagerly():
              if api.is_dtensor(initial_value):
                  value_layout = api.fetch_layout(initial_value)
                  if layout is not None and layout != value_layout:
                      raise errors_impl.InvalidArgumentError(
                          None,
                          None,
                          "Conflicting layout are provided for initial "
                          f"value layout ({value_layout}) and variable ({layout}).",
                      )
                  layout = value_layout
              elif layout is not None:
                  initial_value = api.relayout(initial_value, layout)
              else:
                  raise errors_impl.InvalidArgumentError(
                      None,
                      None,
                      "Neither layout nor DTensor initial value are provided.",
                  )
              self.layout = layout
              with api.default_mesh(layout.mesh):
                  super(DVariable, self).__init__(
                      initial_value, *args, dtype=dtype, **kwargs
                  )
          else:
              # FIXME(175928457): Record value layout in graph mode.
              if layout is not None:
                  initial_value = api.relayout(initial_value, layout)
              super(DVariable, self).__init__(
                  initial_value, *args, dtype=dtype, **kwargs
              )
  
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 1 to create a new DVariable object from inherited base classes: resource_variable_ops with example: obj = DVariable(initial_value)