text: # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Various classes representing distributed values."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import weakref
  
  from tensorflow.python.distribute import device_util
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import distribution_strategy_context as ds_context
  from tensorflow.python.distribute import reduce_util
  from tensorflow.python.eager import context
  from tensorflow.python.eager import tape
  from tensorflow.python.framework import composite_tensor
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.framework import type_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_scope as vs
  from tensorflow.python.ops import variables as variables_lib
  from tensorflow.python.training.saving import saveable_object
  from tensorflow.python.training.saving import saveable_object_util
  from tensorflow.python.training.tracking import base as trackable
  from tensorflow.python.util import nest
  from tensorflow.python.util.tf_export import tf_export
  
  
  def _get_current_replica_id_as_int():
    """Returns the current replica ID as an integer, or `None`."""
    replica_context = ds_context.get_replica_context()
    if replica_context:
      replica_id = replica_context.replica_id_in_sync_group
      if not isinstance(replica_id, int):
        replica_id = tensor_util.constant_value(replica_id)
    else:
      replica_id = distribute_lib.get_update_replica_id()
    return replica_id
  
  
  @tf_export("distribute.DistributedValues", v1=[])
  class DistributedValues(object):
    """Base class for representing distributed values.
  
    A subclass instance of DistributedValues is created when creating variables
    within a distribution strategy, iterating a `tf.Dataset` or through
    `strategy.run`.  This base class should never be instantiated
    directly.  DistributedValues contains a value per replica.  Depending on
    the subclass, the values could either be synced on update, synced on demand,
    or never synced.
  
    DistributedValues can be reduced to obtain single value across replicas,
    as input into `run` or the per replica values inspected
    using `experimental_local_results`.
  
    Example usage:
  
    1. Created from Dataset:
  
    >>> strategy = tf.distribute.MirroredStrategy()
    >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
    >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
    >>> distributed_values = next(dataset_iterator)
  
    2. Returned by `run`:
  
    >>> strategy = tf.distribute.MirroredStrategy()
    >>> @tf.function
    ... def run():
    ...   ctx = tf.distribute.get_replica_context()
    ...   return ctx.replica_id_in_sync_group
    >>> distributed_values = strategy.run(run)
  
    3. As input into `run`:
  
    >>> strategy = tf.distribute.MirroredStrategy()
    >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
    >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
    >>> distributed_values = next(dataset_iterator)
    >>> @tf.function
    ... def run(input):
    ...   return input + 1.0
    >>> updated_value = strategy.run(run, args=(distributed_values,))
  
    4. Reduce value:
  
    >>> strategy = tf.distribute.MirroredStrategy()
    >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
    >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
    >>> distributed_values = next(dataset_iterator)
    >>> reduced_value = strategy.reduce(tf.distribute.ReduceOp.SUM,
    ...                                 distributed_values,
    ...                                 axis = 0)
  
    5. Inspect per replica values:
  
    >>> strategy = tf.distribute.MirroredStrategy()
    >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
    >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
    >>> per_replica_values = strategy.experimental_local_results(
    ...    distributed_values)
    >>> per_replica_values
    (<tf.Tensor: shape=(2,), dtype=float32,
     numpy=array([5., 6.], dtype=float32)>,)
  
    """
  
    def __init__(self, values):
      """Should only be called by subclass __init__."""
      self._values = tuple(values)
  
    def _get(self):
      """Returns the value for the current device or raises a ValueError."""
      replica_id = _get_current_replica_id_as_int()
      if replica_id is None:
        return self._get_cross_replica()
      else:
        return self._values[replica_id]
  
    def _get_cross_replica(self):
      raise NotImplementedError(
          "This method should be overridden by sub-classes which support cross-"
          "replica accesses.")
  
    def _get_closest(self):
      """Returns value in same replica or device if possible, else the _primary."""
      replica_id = _get_current_replica_id_as_int()
      if replica_id is None:
        # Try to find a value on the current device.
        current_device = device_util.canonicalize(device_util.current())
        for value in self._values:
          if device_util.canonicalize(value.device) == current_device:
            return value
        return self._primary
      else:
        return self._values[replica_id]
  
    @property
    def _primary(self):
      """Returns a representative component."""
      return self._values[0]
  
    @property
    def _devices(self):
      return tuple(v.device for v in self._values)
  
    def __str__(self):
      debug_str = ",\n".join(
          "  %d: %s" % (i, v) for i, v in enumerate(self._values))
      return "%s:{\n%s\n}" % (self.__class__.__name__, debug_str)
  
    def __repr__(self):
      debug_repr = ",\n".join(
          "  %d: %r" % (i, v) for i, v in enumerate(self._values))
      return "%s:{\n%s\n}" % (self.__class__.__name__, debug_repr)
  
  
  # NOTE(josh11b,apassos): It would be great if we could inspect the values this was
  # initialized with and use that to generate the overloaded operators here.
  # Unfortunately, Python's rules for special methods don't allow this, see
  # https://docs.python.org/3/reference/datamodel.html#special-method-names
  # "if a class defines a method named __getitem__(), and x is an instance of
  # this class, then x[i] is roughly equivalent to type(x).__getitem__(x, i)."
  # In particular, these special methods don't go through __getattr__, and
  # it will only use those methods if they are defined in the class, not the
  # object.
  class DistributedDelegate(DistributedValues):
    """A map from device to values; acts as the same type as the values."""
  
    def __getattr__(self, name):
      # The '_use_resource_variables' and the attrs starts with '_self' are used
      # for restoring the saved_model proto, and '_attribute_sentinel' is used for
      # Layer tracking. At the point these attrs are queried, the variable has not
      # been initialized. Thus it should not query those of the underlying
      # components.
      if name.startswith("_self_") or name in ("_use_resource_variables",
                                               "_attribute_sentinel",
                                               "_distributed_container"):
        return super(DistributedDelegate, self).__getattr__(name)
  
      # TODO(priyag): This needs to be made robust against pitfalls from mix use
      # __getattr__ and @property. See b/120402273.
      return getattr(self._get(), name)
  
    @property
    def values(self):
      """Returns the per replica values."""
      return self._values
  
    def _get_as_operand(self):
      """Returns the value for operations for the current device.
  
      Some implementations, e.g. `TPUMirroredVariable`, are not able to return the
      value type within a replica context. They can, however, return a value that
      can be used by the operations below.
      """
      return self._get()
  
    # pylint: disable=multiple-statements
    def __add__(self, o):
      return self._get_as_operand() + o
  
    def __radd__(self, o):
      return o + self._get_as_operand()
  
    def __sub__(self, o):
      return self._get_as_operand() - o
  
    def __rsub__(self, o):
      return o - self._get_as_operand()
  
    def __mul__(self, o):
      return self._get_as_operand() * o
  
    def __rmul__(self, o):
      return o * self._get_as_operand()
  
    def __truediv__(self, o):
      return self._get_as_operand() / o
  
    def __rtruediv__(self, o):
      return o / self._get_as_operand()
  
    def __floordiv__(self, o):
      return self._get_as_operand() // o
  
    def __rfloordiv__(self, o):
      return o // self._get_as_operand()
  
    def __mod__(self, o):
      return self._get_as_operand() % o
  
    def __rmod__(self, o):
      return o % self._get_as_operand()
  
    def __lt__(self, o):
      return self._get_as_operand() < o
  
    def __le__(self, o):
      return self._get_as_operand() <= o
  
    def __gt__(self, o):
      return self._get_as_operand() > o
  
    def __ge__(self, o):
      return self._get_as_operand() >= o
  
    def __and__(self, o):
      return self._get_as_operand() & o
  
    def __rand__(self, o):
      return o & self._get_as_operand()
  
    def __or__(self, o):
      return self._get_as_operand() | o
  
    def __ror__(self, o):
      return o | self._get_as_operand()
  
    def __xor__(self, o):
      return self._get_as_operand() ^ o
  
    def __rxor__(self, o):
      return o ^ self._get_as_operand()
  
    def __getitem__(self, o):
      return self._get_as_operand()[o]
  
    def __pow__(self, o, modulo=None):
      return pow(self._get_as_operand(), o, modulo)
  
    def __rpow__(self, o):
      return pow(o, self._get_as_operand())
  
    def __invert__(self):
      return ~self._get_as_operand()
  
    def __neg__(self):
      return -self._get_as_operand()
  
    def __abs__(self):
      return abs(self._get_as_operand())
  
    def __div__(self, o):
      try:
        return self._get_as_operand().__div__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    def __rdiv__(self, o):
      try:
        return self._get_as_operand().__rdiv__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    def __matmul__(self, o):
      try:
        return self._get_as_operand().__matmul__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    def __rmatmul__(self, o):
      try:
        return self._get_as_operand().__rmatmul__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    # TODO(josh11b): Even more operator overloads.
  
  
  class PerReplica(DistributedValues, composite_tensor.CompositeTensor):
    """Holds a map from replica to unsynchronized values."""
  
    @property
    def _type_spec(self):
      return PerReplicaSpec(
          *(type_spec.type_spec_from_value(v) for v in self._values))
  
    @property
    def values(self):
      """Returns the per replica values."""
      return self._values
  
  
  class PerReplicaSpec(type_spec.TypeSpec):
    """Type specification for a `PerReplica`."""
  
    __slots__ = ["_value_specs"]
  
    value_type = property(lambda self: PerReplica)
  
    def __init__(self, *value_specs):
      self._value_specs = tuple(value_specs)
  
    def _serialize(self):
      return self._value_specs
  
    @property
    def _component_specs(self):
      return self._value_specs
  
    def _to_components(self, value):
      replica_context = ds_context.get_replica_context()
      if replica_context is not None and replica_context.num_replicas_in_sync > 1:
        raise ValueError(
            "Flattening a PerReplica to components is not supported in replica "
            "context.")
      return value._values  # pylint: disable=protected-access
  
    def _from_components(self, tensor_list):
      return PerReplica(tensor_list)
  
  
  # Note that unlike PerReplica, Mirrored values inherit from
  # DistributedDelegate and so can be used directly in cross-replica mode.
  # TODO(tomhennigan) Should this extend CompositeTensor?
  class Mirrored(DistributedDelegate):
    """Holds a map from replica to values which are kept in sync."""
  
    def _get_cross_replica(self):
      return self._get_closest()
  
    def _as_graph_element(self):
      obj = self._get()
      conv_fn = getattr(obj, "_as_graph_element", None)
      if conv_fn and callable(conv_fn):
        return conv_fn()
      return obj
  
  
  def _assign_on_device(device, variable, tensor):
    with ops.device(device):
      return variable.assign(tensor)
  
  
  def _assign_add_on_device(device, variable, tensor):
    with ops.device(device):
      return variable.assign_add(tensor)
  
  
  def _assign_sub_on_device(device, variable, tensor):
    with ops.device(device):
      return variable.assign_sub(tensor)
  
  
  class DistributedVarOp(object):
    """A class that looks like `tf.Operation`."""
  
    def __init__(self, name, graph, traceback, typ):
      self.name = name
      self.graph = graph
      self.traceback = traceback
      self.type = typ
  
    def __eq__(self, o):
      if not isinstance(o, self.__class__):
        raise NotImplementedError
      return (self.name == o.name and self.graph == o.graph and
              self.traceback == o.traceback and self.type == o.type)
  
    def __hash__(self):
      return hash((self.name, self.graph, self.traceback, self.type))
  
  
  class DistributedVariable(DistributedDelegate, variables_lib.Variable):
    """Holds a map from replica to variables."""
  
    # TODO(josh11b): Support changing the set of variables if e.g. if new
    # devices are joining or a device is to leave.
  
    def __init__(self, strategy, values, aggregation):
      self._distribute_strategy = strategy
      self._aggregation = aggregation
      super(DistributedVariable, self).__init__(values)
      self._common_name = self._primary.name.split(":")[0]
      # tf.keras keeps track of variables initialized using this attribute. When
      # tf.keras gets the default session, it initializes all uninitialized vars.
      # We need to make _keras_initialized a member of DistributedVariable because
      # without this it will use `__getattr__` which will delegate to a component
      # variable.
      self._keras_initialized = False
      # Typically, a `DistributedVariable`'s initializer is composed of the
      # initializers of the components variables. However, in some cases, such as
      # when restoring from a checkpoint, we may set the _initializer_op
      # property on the entire `DistributedVariable`.
      self._initializer_op = None
  
    def is_initialized(self, name=None):
      """Identifies if all the component variables are initialized.
  
      Args:
        name: Name of the final `logical_and` op.
  
      Returns:
        The op that evaluates to True or False depending on if all the
        component variables are initialized.
      """
      result = self._primary.is_initialized()
      # We iterate through the list of values except the last one to allow us to
      # name the final `logical_and` op the same name that is passed by the user
      # to the `is_initialized` op. For distributed variables, the
      # `is_initialized` op is a `logical_and` op.
      for v in self._values[1:-1]:
        result = math_ops.logical_and(result, v.is_initialized())
      result = math_ops.logical_and(
          result, self._values[-1].is_initialized(), name=name)
      return result
  
    @property
    def initializer(self):
      if self._initializer_op:
        init_op = self._initializer_op
      else:
        # return grouped ops of all the var initializations of component values of
        # the mirrored variable
        init_op = control_flow_ops.group(
            tuple(v.initializer for v in self._values))
      return init_op
  
    def initialized_value(self):
      return self._get_closest().initialized_value()
  
    @property
    def initial_value(self):
      return self._get_closest().initial_value
  
    @property
    def constraint(self):
      return self._primary.constraint
  
    @property
    def graph(self):
      return self._primary.graph
  
    @property
    def _shared_name(self):
      return self._common_name
  
    @property
    def _unique_id(self):
      return self._primary._unique_id  # pylint: disable=protected-access
  
    @property
    def _graph_key(self):
      """Lets Optimizers know which graph this variable is from."""
      return self._primary._graph_key  # pylint: disable=protected-access
  
    @property
    def name(self):
      return self._primary.name
  
    @property
    def dtype(self):
      return self._primary.dtype
  
    @property
    def shape(self):
      return self._primary.shape
  
    @property
    def synchronization(self):
      return self._primary.synchronization
  
    @property
    def aggregation(self):
      return self._aggregation
  
    @property
    def handle(self):
      replica_id = _get_current_replica_id_as_int()
      if replica_id is None:
        raise ValueError("`handle` is not available outside the replica context"
                         " or a `tf.distribute.Strategy.update()` call.")
      else:
        return self._values[replica_id].handle
  
    def eval(self, session=None):
      return self._get_closest().eval(session)
  
    @property
    def _save_slice_info(self):
      return self._primary._save_slice_info  # pylint: disable=protected-access
  
    def _get_save_slice_info(self):
      return self._primary._get_save_slice_info()  # pylint: disable=protected-access
  
    def _set_save_slice_info(self, save_slice_info):
      for v in self._values:
        v._set_save_slice_info(save_slice_info)  # pylint: disable=protected-access
  
    @property
    def device(self):
      return self._get_closest().device
  
    @property
    def trainable(self):
      return self._primary.trainable
  
    @property
    def distribute_strategy(self):
      return self._distribute_strategy
  
    def get_shape(self):
      return self._primary.get_shape()
  
    def to_proto(self, export_scope=None):
      return self._primary.to_proto(export_scope=export_scope)
  
    @property
    def op(self):
      # We want cross-replica code that does some var.op.X calls
      # to work (even if the current device isn't in self._devices), but
      # other uses of var.op in a cross-replica context to fail.
      if ds_context.in_cross_replica_context():
        return DistributedVarOp(self._primary.op.name, self._primary.op.graph,
                                self._primary.op.traceback, self._primary.op.type)
      return self._get().op
  
    @property
    def _in_graph_mode(self):
      return self._primary._in_graph_mode  # pylint: disable=protected-access
  
    def read_value(self):
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        return array_ops.identity(self._get())
  
    def value(self):
      return self._get_closest().value()
  
    def assign_sub(self, value, use_locking=False, name=None, read_value=True):
      assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw)
      return self._update(
          update_fn=assign_sub_fn,
          value=value,
          use_locking=use_locking,
          name=name,
          read_value=read_value)
  
    def assign_add(self, value, use_locking=False, name=None, read_value=True):
      assign_add_fn = lambda var, *a, **kw: var.assign_add(*a, **kw)
      return self._update(
          update_fn=assign_add_fn,
          value=value,
          use_locking=use_locking,
          name=name,
          read_value=read_value)
  
    def assign(self, value, use_locking=False, name=None, read_value=True):
      assign_fn = lambda var, *a, **kw: var.assign(*a, **kw)
      return self._update(
          update_fn=assign_fn,
          value=value,
          use_locking=use_locking,
          name=name,
          read_value=read_value)
  
    def scatter_sub(self, sparse_delta, use_locking=False, name=None):
      scatter_sub_fn = lambda var, *a, **kw: var.scatter_sub(*a, **kw)
      return self._update(
          update_fn=scatter_sub_fn,
          value=sparse_delta,
          use_locking=use_locking,
          name=name)
  
    def scatter_add(self, sparse_delta, use_locking=False, name=None):
      scatter_add_fn = lambda var, *a, **kw: var.scatter_add(*a, **kw)
      return self._update(
          update_fn=scatter_add_fn,
          value=sparse_delta,
          use_locking=use_locking,
          name=name)
  
    def scatter_mul(self, sparse_delta, use_locking=False, name=None):
      scatter_mul_fn = lambda var, *a, **kw: var.scatter_mul(*a, **kw)
      return self._update(
          update_fn=scatter_mul_fn,
          value=sparse_delta,
          use_locking=use_locking,
          name=name)
  
    def scatter_div(self, sparse_delta, use_locking=False, name=None):
      scatter_div_fn = lambda var, *a, **kw: var.scatter_div(*a, **kw)
      return self._update(
          update_fn=scatter_div_fn,
          value=sparse_delta,
          use_locking=use_locking,
          name=name)
  
    def scatter_min(self, sparse_delta, use_locking=False, name=None):
      scatter_min_fn = lambda var, *a, **kw: var.scatter_min(*a, **kw)
      return self._update(
          update_fn=scatter_min_fn,
          value=sparse_delta,
          use_locking=use_locking,
          name=name)
  
    def scatter_max(self, sparse_delta, use_locking=False, name=None):
      scatter_max_fn = lambda var, *a, **kw: var.scatter_max(*a, **kw)
      return self._update(
          update_fn=scatter_max_fn,
          value=sparse_delta,
          use_locking=use_locking,
          name=name)
  
    def scatter_update(self, sparse_delta, use_locking=False, name=None):
      scatter_update_fn = lambda var, *a, **kw: var.scatter_update(*a, **kw)
      return self._update(
          update_fn=scatter_update_fn,
          value=sparse_delta,
          use_locking=use_locking,
          name=name)
  
    def _update_cross_replica(self, update_fn, value, **kwargs):
      """Applies updates across replicas.
  
      Args:
        update_fn: A callable to pass to `strategy.extended.update` to update the
          variable. It should has the same signature as `Variable.assign()`.
        value: value to be passed to `update_fn`.
        **kwargs: remaining arguments to `update_fn`.
  
      Returns:
        Updated variable or `tf.Operation`.
      """
      return self.distribute_strategy.extended.update(
          self, update_fn, args=(value,), kwargs=kwargs, group=True)
  
    def _update_replica(self, update_fn, value, **kwargs):
      """Applies updates in one replica.
  
      Args:
        update_fn: A callable to update the variable. It should has the same
          signature as `Variable.assign()`.
        value: value to be passed to `update_fn`.
        **kwargs: remaining arguments to `update_fn`.
  
      Returns:
        Updated variable or `tf.Operation`.
      """
      raise NotImplementedError("should be implemented by subclass.")
  
    def _update(self, update_fn, value, **kwargs):
      """Applies updates depending on the context.
  
      The method calls `_update_replica` in replica context,
      `_update_cross_replica` in cross replica context, and `update_fn` in update
      context.
  
      If `read_value` is True, the method returns the updated Variable. If
      `read_value` is False, the method returns the update `tf.Operation`.
  
      Args:
        update_fn: A callable to pass to `strategy.extended.update` to update the
          variable. It should have the same signature as `Variable.assign()`.
        value: value to be passed to `update_fn`.
        **kwargs: keyword arguments to `update_fn`.
  
      Returns:
        Updated variable or `tf.Operation`.
  
      """
      with ds_context.enter_or_assert_strategy(self.distribute_strategy):
        if ds_context.in_cross_replica_context():
          update_replica_id = distribute_lib.get_update_replica_id()
          if update_replica_id is not None:
            return update_fn(self._values[update_replica_id], value, **kwargs)
          return self._update_cross_replica(update_fn, value, **kwargs)
        else:
          _assert_replica_context(self.distribute_strategy)
          return self._update_replica(update_fn, value, **kwargs)
  
    def _should_act_as_resource_variable(self):
      """Pass resource_variable_ops.is_resource_variable check."""
      pass
  
  
  ops.register_dense_tensor_like_type(DistributedVariable)
  
  
  def _validate_colocate_extended(v, extended):
    variable_strategy = v._distribute_strategy  # pylint: disable=protected-access
    if variable_strategy.extended is not extended:
      raise ValueError(
          "`colocate_vars_with` must only be passed a variable created in this "
          "tf.distribute.Strategy.scope(), not %s created in scope: %s" %
          (v, variable_strategy))
  
  
  def validate_colocate_distributed_variable(v, extended):
    if not isinstance(v, DistributedVariable):
      raise ValueError(
          "`colocate_vars_with` must only be passed a variable created in this "
          "tf.distribute.Strategy.scope(), not: %r" % (v,))
    _validate_colocate_extended(v, extended)
  
  
  def validate_colocate(v, extended):
    if not hasattr(v, "_distribute_strategy"):
      raise ValueError(
          "`colocate_vars_with` must only be passed a variable created in this "
          "tf.distribute.Strategy.scope(), not: %r" % (v,))
    _validate_colocate_extended(v, extended)
  
  
  def _apply_aggregation(strategy, value, aggregation, destinations):
    if aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
      return strategy.extended.broadcast_to(
          strategy.experimental_local_results(value)[0],
          destinations=destinations)
    reduce_op = reduce_util.ReduceOp.from_variable_aggregation(aggregation)
    return strategy.extended.reduce_to(reduce_op, value, destinations)
  
  
  _aggregation_error_msg = (
      "You must specify an aggregation method to update a "
      "{variable_type} in Replica Context. You can do so by passing "
      "an explicit value for argument `aggregation` to tf.Variable(..)."
      "e.g. `tf.Variable(..., aggregation=tf.VariableAggregation.SUM)`"
      "`tf.VariableAggregation` lists the possible aggregation methods."
      "This is required because {variable_type} should always be "
      "kept in sync. When updating them or assigning to them in a "
      "replica context, we automatically try to aggregate the values "
      "before updating the variable. For this aggregation, we need to "
      "know the aggregation method. "
      "Another alternative is to not try to update such "
      "{variable_type} in replica context, but in cross replica "
      "context. You can enter cross replica context by calling "
      "`tf.distribute.get_replica_context().merge_call(merge_fn, ..)`."
      "Inside `merge_fn`, you can then update the {variable_type} "
      "using `tf.distribute.StrategyExtended.update()`.")
  
  
  class _MirroredSaveable(saveable_object_util.ResourceVariableSaveable):
    """Class for defining how to restore a MirroredVariable."""
  
    def __init__(self, mirrored_variable, primary_variable, name):
      self._mirrored_variable = mirrored_variable
      super(_MirroredSaveable, self).__init__(primary_variable, "", name)
  
    def restore(self, restored_tensors, restored_shapes):
      """Restore the same value into all variables."""
      tensor, = restored_tensors
      return control_flow_ops.group(
          tuple(
              _assign_on_device(v.device, v, tensor)
              for v in self._mirrored_variable.values))
  
  
  def create_mirrored_variable(  # pylint: disable=missing-docstring
      strategy, real_mirrored_creator, mirrored_cls, sync_on_read_cls, **kwargs):
    # Figure out what collections this variable should be added to.
    # We'll add the MirroredVariable to those collections instead.
    var_collections = kwargs.pop("collections", None)
    if var_collections is None:
      var_collections = [ops.GraphKeys.GLOBAL_VARIABLES]
    kwargs["collections"] = []
  
    synchronization = kwargs.get("synchronization",
                                 vs.VariableSynchronization.ON_WRITE)
  
    if synchronization == vs.VariableSynchronization.NONE:
      raise ValueError(
          "`NONE` variable synchronization mode is not supported with `Mirrored` "
          "distribution strategy. Please change the `synchronization` for "
          "variable: " + str(kwargs["name"]))
    elif synchronization == vs.VariableSynchronization.ON_READ:
      is_sync_on_read = True
    elif synchronization in (vs.VariableSynchronization.ON_WRITE,
                             vs.VariableSynchronization.AUTO):
      # `AUTO` synchronization defaults to `ON_WRITE`.
      is_sync_on_read = False
    else:
      raise ValueError(
          "Invalid variable synchronization mode: %s for variable: %s" %
          (synchronization, kwargs["name"]))
  
    aggregation = kwargs.pop("aggregation", vs.VariableAggregation.NONE)
  
    if aggregation not in (vs.VariableAggregation.NONE,
                           vs.VariableAggregation.SUM,
                           vs.VariableAggregation.MEAN,
                           vs.VariableAggregation.ONLY_FIRST_REPLICA):
      raise ValueError("Invalid variable aggregation mode: %s for variable: %s" %
                       (aggregation, kwargs["name"]))
  
    # Ignore user-specified caching device, not needed for mirrored variables.
    kwargs.pop("caching_device", None)
  
    # TODO(josh11b,apassos): It would be better if variable initialization
    # was never recorded on the tape instead of having to do this manually
    # here.
    with tape.stop_recording():
      value_list = real_mirrored_creator(**kwargs)
      var_cls = sync_on_read_cls if is_sync_on_read else mirrored_cls
      result = var_cls(strategy, value_list, aggregation)
      # Install the created DistributedVariable as _distributed_container property
      # of the underlying variables, to make it easy to map back to the container.
      for v in result.values:
        # Hold a strong reference to avoid the container from being GC-ed. After
        # v = v.assign(), the user code may no longer holds references to the
        # original container, since v.assign() returns a new DistributedVariable.
        v._distributed_container = result  # pylint: disable=protected-access
  
    # Add the wrapped variable to the requested collections.
    # The handling of eager mode and the global step matches
    # ResourceVariable._init_from_args().
    if not context.executing_eagerly():
      g = ops.get_default_graph()
      # If "trainable" is True, next_creator() will add the member variables
      # to the TRAINABLE_VARIABLES collection, so we manually remove
      # them and replace with the MirroredVariable. We can't set
      # "trainable" to False for next_creator() since that causes functions
      # like implicit_gradients to skip those variables.
      if kwargs.get("trainable", True):
        var_collections.append(ops.GraphKeys.TRAINABLE_VARIABLES)
        l = g.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES)
        for value in value_list:
          for i, trainable_variable in enumerate(l):
            if value is trainable_variable:
              del l[i]
              break
  
      g.add_to_collections(var_collections, result)
    elif ops.GraphKeys.GLOBAL_STEP in var_collections:
      ops.add_to_collections(ops.GraphKeys.GLOBAL_STEP, result)
  
    return result
  
  
  class MirroredVariable(DistributedVariable, Mirrored):
    """Holds a map from replica to variables whose values are kept in sync."""
  
    def _update_replica(self, update_fn, value, **kwargs):
      if self.aggregation == vs.VariableAggregation.NONE:
        raise ValueError(
            _aggregation_error_msg.format(variable_type="MirroredVariable"))
  
      def merge_fn(strategy, value, **kwargs):
        """Aggregate values and update all variables in cross replica context."""
        # Don't allow MEAN with non float dtype, since it may cause unexpected
        # precision loss. Python3 and NumPy automatically upcast integers to
        # float in division, but we should always preserve the type.
        #
        # Note that to be backward compatible we allow the case when the value
        # is *always* the same on each replica. I.E. value is not a
        # PerReplica. Refer to regroup() to see how values are grouped.
        if self._aggregation == vs.VariableAggregation.MEAN and (
            not self.dtype.is_floating) and isinstance(value, PerReplica):
          raise ValueError(
              "Cannot update non-float variables with "
              "tf.VariableAggregation.MEAN aggregation in replica context. "
              "Either change the variable dtype to float or update it in "
              "cross-replica context.")
  
        assert strategy == self.distribute_strategy
        v = _apply_aggregation(strategy, value, self.aggregation, self)
        return self._update_cross_replica(update_fn, v, **kwargs)
  
      return ds_context.get_replica_context().merge_call(
          merge_fn, args=(value,), kwargs=kwargs)
  
    def scatter_min(self, *args, **kwargs):
      if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
          self._aggregation != vs.VariableAggregation.NONE):
        raise NotImplementedError("scatter_min is only supported for mirrored "
                                  "variable (variable created within certain "
                                  "`tf.distribute.Strategy` scope) with NONE or "
                                  "`ONLY_FIRST_REPLICA` aggregation, got: %s" %
                                  self._aggregation)
      return super(MirroredVariable, self).scatter_min(*args, **kwargs)
  
    def scatter_max(self, *args, **kwargs):
      if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
          self._aggregation != vs.VariableAggregation.NONE):
        raise NotImplementedError("scatter_max is only supported for mirrored "
                                  "variable (variable created within certain "
                                  "`tf.distribute.Strategy` scope) with NONE or "
                                  "`ONLY_FIRST_REPLICA` aggregation, got: %s" %
                                  self._aggregation)
      return super(MirroredVariable, self).scatter_max(*args, **kwargs)
  
    def scatter_update(self, *args, **kwargs):
      if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
          self._aggregation != vs.VariableAggregation.NONE):
        raise NotImplementedError("scatter_update is only supported for mirrored "
                                  "variable (variable created within certain "
                                  "`tf.distribute.Strategy` scope) with NONE or "
                                  "`ONLY_FIRST_REPLICA` aggregation, got: %s" %
                                  self._aggregation)
      return super(MirroredVariable, self).scatter_update(*args, **kwargs)
  
    def _get_cross_replica(self):
      # Return identity, to avoid directly exposing the variable to the user and
      # allowing it to be modified by mistake.
      return array_ops.identity(Mirrored._get_cross_replica(self))
  
    def _as_graph_element(self):
      return self._get_closest()._as_graph_element()  # pylint: disable=protected-access
  
    def _gather_saveables_for_checkpoint(self):
      """Overrides Trackable method.
  
      This allows both name-based and object-based save and restore of
      MirroredVariables.
  
      Returns:
        A dictionary mapping attribute names to `SaveableObject` factories.
      """
  
      def _saveable_factory(name=self._common_name):
        return _MirroredSaveable(self, self._primary, name)
  
      return {trackable.VARIABLE_VALUE_KEY: _saveable_factory}
  
    def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
      """Converts a variable to a tensor."""
      # Try to avoid assignments to and other mutations of MirroredVariable
      # state except through a DistributionStrategy.extended.update() call.
      if as_ref:
        # A TF 1.x case where the variable is a boolean variable and used like:
        # tf.cond(v, true_fn, false_fn).
        raise ValueError(
            "You may be using variable created under distribute strategy in TF "
            "1.x control flows. Try explicitly converting the variable to Tensor "
            "using variable.read_value(), or switch to TF 2.x.")
      return ops.convert_to_tensor(
          self._get(), dtype=dtype, name=name, as_ref=as_ref)
  
  
  # Register a conversion function which reads the value of the variable,
  # allowing instances of the class to be used as tensors.
  def _tensor_conversion_mirrored(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
  
  
  ops.register_tensor_conversion_function(MirroredVariable,
                                          _tensor_conversion_mirrored)
  
  
  def _tensor_conversion_mirrored_val(value, dtype=None, name=None, as_ref=False):
    return ops.convert_to_tensor(
        value._get(), dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access
  
  
  ops.register_tensor_conversion_function(Mirrored,
                                          _tensor_conversion_mirrored_val)
  
  
  def is_distributed_variable(v):
    """Determine if a variable is ds variable or TPU mirrored variable."""
    return isinstance(v, DistributedVariable)
  
  
  class _SyncOnReadSaveable(saveable_object.SaveableObject):
    """Class for defining how to restore a SyncOnReadVariable."""
  
    def __init__(self, sync_on_read_variable, name):
      self._sync_on_read_variable = sync_on_read_variable
  
      # We use a callable so that we don't have to evaluate this expression
      # in the case where we are trying to restore instead of save.
      def tensor():
        strategy = sync_on_read_variable._distribute_strategy  # pylint: disable=protected-access
        return strategy.extended.read_var(sync_on_read_variable)
  
      spec = saveable_object.SaveSpec(
          tensor=tensor,
          slice_spec="",
          name=name,
          dtype=sync_on_read_variable.dtype,
          device=sync_on_read_variable._primary.device)  # pylint: disable=protected-access
  
      super(_SyncOnReadSaveable, self).__init__(tensor, [spec], name)
  
    def restore(self, restored_tensors, restored_shapes):
      """Restore the same value into all variables."""
      # To preserve the sum across save and restore, we have to divide the
      # total across all devices when restoring a variable that was summed
      # when saving.
      tensor, = restored_tensors
      if self._sync_on_read_variable.aggregation == vs.VariableAggregation.SUM:
        tensor = math_ops.cast(tensor / len(self._sync_on_read_variable._devices),  # pylint: disable=protected-access
                               self._sync_on_read_variable.dtype)
      return control_flow_ops.group(
          tuple(
              _assign_on_device(v.device, v, tensor)
              for v in self._sync_on_read_variable.values))
  
  
  def _assert_replica_context(strategy):
    replica_context = ds_context.get_replica_context()
    if not replica_context:
      raise RuntimeError(
          "Replica-local variables may only be assigned in a replica context.")
    if replica_context.strategy is not strategy:
      raise RuntimeError(
          "Replica-local variables may only be assigned in a replica context.")
  
  
  class SyncOnReadVariable(DistributedVariable):
    """Holds a map from replica to variables whose values are reduced on save."""
  
    def _update_replica(self, update_fn, value, **kwargs):
      return update_fn(self._get_closest(), value, **kwargs)
  
    # TODO(b/154017756): Make assign behaivor in cross replica context consistent
    # with MirroredVariable.
    def assign_sub(self, *args, **kwargs):
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        if ds_context.in_cross_replica_context():
          if self._aggregation == vs.VariableAggregation.SUM:
            raise ValueError(
                "SyncOnReadVariable does not support `assign_sub` in "
                "cross-replica context when aggregation is set to "
                "`tf.VariableAggregation.SUM`.")
          return control_flow_ops.group(
              tuple(
                  _assign_sub_on_device(v.device, v, args[0])
                  for v in self._values))
        else:
          return super(SyncOnReadVariable, self).assign_sub(*args, **kwargs)
  
    def assign_add(self, *args, **kwargs):
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        if ds_context.in_cross_replica_context():
          if self._aggregation == vs.VariableAggregation.SUM:
            raise ValueError(
                "SyncOnReadVariable does not support `assign_add` in "
                "cross-replica context when aggregation is set to "
                "`tf.VariableAggregation.SUM`.")
          return control_flow_ops.group(
              tuple(
                  _assign_add_on_device(v.device, v, args[0])
                  for v in self._values))
        else:
          return super(SyncOnReadVariable, self).assign_add(*args, **kwargs)
  
    def assign(self, *args, **kwargs):
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        if ds_context.in_cross_replica_context():
          # To preserve the sum across save and restore, we have to divide the
          # total across all devices when restoring a variable that was summed
          # when saving.
          tensor = args[0]
          if self._aggregation == vs.VariableAggregation.SUM:
            tensor = math_ops.cast(tensor / len(self._values), self.dtype)
          return control_flow_ops.group(
              tuple(_assign_on_device(v.device, v, tensor) for v in self._values))
        else:
          return super(SyncOnReadVariable, self).assign(*args, **kwargs)
  
    def _scatter_not_implemented(self, method):
      raise NotImplementedError(
          "Variables with `synchronization=ON_READ` doesn't support `%s`" %
          method)
  
    def scatter_sub(self, *args, **kwargs):
      self._scatter_not_implemented("scatter_sub")
  
    def scatter_add(self, *args, **kwargs):
      self._scatter_not_implemented("scatter_add")
  
    def scatter_mul(self, *args, **kwargs):
      self._scatter_not_implemented("scatter_mul")
  
    def scatter_div(self, *args, **kwargs):
      self._scatter_not_implemented("scatter_div")
  
    def scatter_min(self, *args, **kwargs):
      self._scatter_not_implemented("scatter_min")
  
    def scatter_max(self, *args, **kwargs):
      self._scatter_not_implemented("scatter_max")
  
    def scatter_update(self, *args, **kwargs):
      self._scatter_not_implemented("scatter_update")
  
    def value(self):
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        if ds_context.in_cross_replica_context():
          return self._get_cross_replica()
        else:
          # _get_closest() returns a Variable.
          return self._get_closest().value()
  
    def numpy(self):
      if context.executing_eagerly():
        return self.read_value().numpy()
      else:
        raise NotImplementedError(
            "numpy() is only available when eager execution is enabled.")
  
    def _get_cross_replica(self):
      if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
        return self._primary
  
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        return self._distribute_strategy.reduce(
            reduce_util.ReduceOp.from_variable_aggregation(self.aggregation),
            self,
            axis=None)
  
    def _as_graph_element(self):
      # pylint: disable=protected-access
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        if ds_context.in_cross_replica_context():
          return ops.convert_to_tensor(self._get_cross_replica())
      return self._get()._as_graph_element()
  
    def _gather_saveables_for_checkpoint(self):
      """Overrides Trackable method.
  
      This allows both name-based and object-based save and restore of
      `SyncOnReadVariable`s.
  
      Returns:
        A dictionary mapping attribute names to `SaveableObject` factories.
      """
  
      def _saveable_factory(name=self._common_name):
        return _SyncOnReadSaveable(self, name)
  
      return {trackable.VARIABLE_VALUE_KEY: _saveable_factory}
  
    def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
      """Converts a variable to a tensor."""
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        return ops.convert_to_tensor(
            self._get(), dtype=dtype, name=name, as_ref=as_ref)
  
  
  # Register a conversion function for SyncOnReadVariable which allows as_ref to
  # be true.
  def _tensor_conversion_sync_on_read(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
  
  
  ops.register_tensor_conversion_function(SyncOnReadVariable,
                                          _tensor_conversion_sync_on_read)
  
  
  def regroup(values, wrap_class=PerReplica, always_wrap=False):
    """Makes a nest per-replica into a nest of PerReplica/Mirrored values.
  
    Args:
      values: Values to regroup
      wrap_class: Class that `values` be wrapped in.
      always_wrap: Always wrap the `values` in `wrap_class` even if the values
          are the same except for DistributeVariable.
    Returns:
      Wrapped `values`.
    """
    v0 = values[0]
  
    if isinstance(v0, list):
      for v in values[1:]:
        assert isinstance(v, list)
        assert len(v) == len(v0), ("len(v) == %d, len(v0) == %d, v: %s, v0: %s" %
                                   (len(v), len(v0), v, v0))
      return [
          regroup(tuple(v[i] for v in values), wrap_class)
          for i in range(len(v0))
      ]
  
    if isinstance(v0, tuple):
      for v in values[1:]:
        assert isinstance(v, tuple)
        assert len(v) == len(v0)
      regrouped_tuple = tuple(
          regroup(tuple(v[i] for v in values), wrap_class)
          for i in range(len(v0)))
      if hasattr(v0, "_fields"):
        # This tuple is in fact a namedtuple! Create a new namedtuple instance
        # and initialize it with the regrouped values:
        assert hasattr(type(v0), "_make")
        return type(v0)._make(regrouped_tuple)
      else:
        return regrouped_tuple
  
    if isinstance(v0, dict):
      v0keys = v0.keys()
      for v in values[1:]:
        assert isinstance(v, dict), ("v[0]: %r  v[i]: %r" % (v0, v))
        assert set(v.keys()) == set(v0keys), ("v[0].keys: %s  v[i].keys: %s" %
                                              (set(v0keys), set(v.keys())))
      # Use the actual type in case it is a class inherited from a dict.
      return type(v0)({
          key: regroup(tuple(v[key] for v in values), wrap_class)
          for key in v0keys
      })
  
    # If exactly the same object across all devices, return it unwrapped.
    same_id = True
    for v in values[1:]:
      if v is not v0:
        same_id = False
        break
    # Consider three cases where same_id is true:
    # * If v0 is a DistributedVariable (a MirroredVariable or
    #   SyncOnReadVariable, and same_id means it is the same across all
    #   devices), we want to return it. We check DistributedVariable
    #   specifically since it can look like it has a
    #   _distributed_container member since its members do.
    if same_id and isinstance(v0, DistributedVariable):
      return v0
    # * If v0 is a member of a distributed variable, in which case
    #   hasattr(v0, "_distributed_container") is true, we want to
    #   return the DistributedVariable that contains it using the
    #   _distributed_container logic below. This case can trigger
    #   same_id when there is only one device.
    # * In any other situation, same_id means we return v0 unless `always_wrap` is
    #   true.
    if same_id and not always_wrap and not hasattr(v0, "_distributed_container"):
      return v0
  
    # Detect the case where each device has a parallel component of the
    # same MirroredVariable (or SyncOnReadVariable). In this case we
    # want to return the containing MirroredVariable, after a bunch of
    # sanity checking. In particular, each component should have the
    # same container, and the devices of the variables should match the
    # keys of the per-replica dictionary.
    if hasattr(v0, "_distributed_container"):
      # pylint: disable=protected-access
      assert not isinstance(v0, MirroredVariable), (
          "ids = %s, values = %s" % ([id(v) for v in values], values))
      distributed_container = v0._distributed_container
      assert distributed_container is not None
      for v in values[1:]:
        assert distributed_container is v._distributed_container
      return distributed_container
    # pylint: enable=protected-access
  
    return wrap_class(values)
  
  
  def select_replica(replica_id, structured):
    """Specialize a nest of regular & per-replica values for one replica."""
  
    def _get(x):
      # `DistributedValues` would be sliced according to replica unless it is a
      # `DistributedVariable` because `DistributedVariable` can be handled
      # directly in the replica context.
      if (isinstance(x, DistributedVariable) or
          not isinstance(x, DistributedValues)):
        return x
      else:
        return x.values[replica_id]
  
    return nest.map_structure(_get, structured)
  
  
  def select_replica_mirrored(replica_id, structured):
    """Specialize a nest of regular & mirrored values for one replica."""
  
    def _get_mirrored(x):
      if isinstance(x, DistributedValues):
        if not isinstance(x, Mirrored):
          raise TypeError(
              "Expected value to be mirrored across replicas: %s in %s." %
              (x, structured))
        return x.values[replica_id]
      else:
        return x
  
    return nest.map_structure(_get_mirrored, structured)
  
  
  def update_regroup(extended, updates, group):
    """Regroup for an update, with dependencies to ensure all updates execute."""
    if not group:
      regrouped = regroup(updates, Mirrored)
      return nest.map_structure(extended._local_results, regrouped)  # pylint: disable=protected-access
  
    def _make_grouped_mirrored(values):
      """Convert per-replica list `values` into Mirrored type with grouping."""
      if len(values) == 1:
        return Mirrored(values)
  
      # Make sure we run all updates. Without this, something like
      # session.run(extended.update(...)) may only update one replica.
      g = control_flow_ops.group(values)
  
      # If values is just ops, the grouping is enough. Everything in values
      # should have the same type, since we expect every replica to be performing
      # the same computation.
      if not all(tensor_util.is_tensor(v) for v in values):
        return g
  
      # Otherwise we need tensors with the same values as `values`, but
      # that have a dependency on `g`.
      with_dep = []
      for v in values:
        with ops.device(v.device), ops.control_dependencies([g]):
          with_dep.append(array_ops.identity(v))
  
      return Mirrored(with_dep)
  
    return regroup(updates, _make_grouped_mirrored)
  
  
  def value_container(val):
    """Returns the container that this per-replica `value` belongs to.
  
    Args:
      val: A value returned by `call_for_each_replica()` or a variable created in
        `scope()`.
  
    Returns:
      A container that `value` belongs to.
      If value does not belong to any container (including the case of
      container having been destroyed), returns the value itself.
    """
    if (hasattr(val, "_distributed_container") and
        # DistributedVariable has _distributed_container defined
        # but we don't want to return it.
        not isinstance(val, DistributedVariable)):
      container = val._distributed_container  # pylint: disable=protected-access
      if container is not None:
        return container
    return val
  
  
  class AggregatingVariable(variables_lib.Variable):
    """A wrapper around a variable that aggregates updates across replicas."""
  
    def __init__(self, strategy, v, aggregation):
      self._distribute_strategy = strategy
      self._v = v
      # NOTE: We don't use "_distributed_container" here because we don't want
      # to trigger that code path in regroup().
      v._aggregating_container = weakref.ref(self)  # pylint: disable=protected-access
      self._aggregation = aggregation
  
    def get(self):
      return self._v
  
    @property
    def distribute_strategy(self):
      return self._distribute_strategy
  
    def __getattr__(self, name):
      return getattr(self._v, name)
  
    def _assign_func(self, *args, **kwargs):
      with ds_context.enter_or_assert_strategy(self._distribute_strategy):
        f = kwargs.pop("f")
        if ds_context.in_cross_replica_context():
          if distribute_lib.get_update_replica_id() is not None:
            # We are calling an assign function in an update context.
            return f(self._v, *args, **kwargs)
  
          # We are calling an assign function in cross replica context, wrap it in
          # an update call.
          return self._distribute_strategy.extended.update(
              self, f, args=args, kwargs=kwargs)
        else:
          replica_context = ds_context.get_replica_context()
          assert replica_context
          # We are calling an assign function in replica context.
          # We reduce the value we want to assign/add/sub. More details about how
          # we handle the different use cases can be found in the _reduce method.
          # We call the function with the reduced value.
          if self._aggregation == vs.VariableAggregation.NONE:
            raise ValueError(
                _aggregation_error_msg.format(
                    variable_type="AggregatingVariable"))
  
          def merge_fn(strategy, value, *other_args, **other_kwargs):
            v = _apply_aggregation(strategy, value, self._aggregation, self)
            return strategy.extended.update(
                self, f, args=(v,) + other_args, kwargs=other_kwargs)
  
          return replica_context.merge_call(merge_fn, args=args, kwargs=kwargs)
  
    def assign_sub(self, *args, **kwargs):
      assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw)
      return self._assign_func(f=assign_sub_fn, *args, **kwargs)
  
    def assign_add(self, *args, **kwargs):
      assign_add_fn = lambda var, *a, **kw: var.assign_add(*a, **kw)
      return self._assign_func(f=assign_add_fn, *args, **kwargs)
  
    def assign(self, *args, **kwargs):
      assign_fn = lambda var, *a, **kw: var.assign(*a, **kw)
      return self._assign_func(f=assign_fn, *args, **kwargs)
  
    @property
    def initializer(self):
      return self._v.initializer
  
    def initialized_value(self):
      return self._v.initialized_value()
  
    @property
    def initial_value(self):
      return self._v.initial_value
  
    @property
    def op(self):
      return self._v.op
  
    def read_value(self):
      return self._v.read_value()
  
    def eval(self, session=None):
      return self._v.eval(session)
  
    @property
    def graph(self):
      return self._v.graph
  
    @property
    def device(self):
      return self._v.device
  
    @property
    def shape(self):
      return self._v.shape
  
    @property
    def aggregation(self):
      return self._aggregation
  
    @property
    def synchronization(self):
      return self._v.synchronization
  
    @property
    def name(self):
      return self._v.name
  
    @property
    def trainable(self):
      return self._v.trainable
  
    @property
    def dtype(self):
      return self._v.dtype
  
    # TODO(josh11b): Test saving & restoring.
    def _gather_saveables_for_checkpoint(self):
      return {trackable.VARIABLE_VALUE_KEY: self._v}
  
    # pylint: disable=multiple-statements
    def __add__(self, o):
      return self._v + o
  
    def __radd__(self, o):
      return o + self._v
  
    def __sub__(self, o):
      return self._v - o
  
    def __rsub__(self, o):
      return o - self._v
  
    def __mul__(self, o):
      return self._v * o
  
    def __rmul__(self, o):
      return o * self._v
  
    def __truediv__(self, o):
      return self._v / o
  
    def __rtruediv__(self, o):
      return o / self._v
  
    def __floordiv__(self, o):
      return self._v // o
  
    def __rfloordiv__(self, o):
      return o // self._v
  
    def __mod__(self, o):
      return self._v % o
  
    def __rmod__(self, o):
      return o % self._v
  
    def __lt__(self, o):
      return self._v < o
  
    def __le__(self, o):
      return self._v <= o
  
    def __gt__(self, o):
      return self._v > o
  
    def __ge__(self, o):
      return self._v >= o
  
    def __and__(self, o):
      return self._v & o
  
    def __rand__(self, o):
      return o & self._v
  
    def __or__(self, o):
      return self._v | o
  
    def __ror__(self, o):
      return o | self._v
  
    def __xor__(self, o):
      return self._v ^ o
  
    def __rxor__(self, o):
      return o ^ self._v
  
    def __getitem__(self, o):
      return self._v[o]
  
    def __pow__(self, o, modulo=None):
      return pow(self._v, o, modulo)
  
    def __rpow__(self, o):
      return pow(o, self._v)
  
    def __invert__(self):
      return ~self._v
  
    def __neg__(self):
      return -self._v
  
    def __abs__(self):
      return abs(self._v)
  
    def __div__(self, o):
      try:
        return self._v.__div__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    def __rdiv__(self, o):
      try:
        return self._v.__rdiv__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    def __matmul__(self, o):
      try:
        return self._v.__matmul__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    def __rmatmul__(self, o):
      try:
        return self._v.__rmatmul__(o)
      except AttributeError:
        # See https://docs.python.org/3/library/constants.html#NotImplemented
        return NotImplemented
  
    def __str__(self):
      return str(self._v)
  
    def __repr__(self):
      return repr(self._v)
  
    def _should_act_as_resource_variable(self):
      """Pass resource_variable_ops.is_resource_variable check."""
      pass
  
    def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
      return ops.convert_to_tensor(self.get(), dtype=dtype, name=name,
                                   as_ref=as_ref)
  
  
  # Register a conversion function which reads the value of the variable,
  # allowing instances of the class to be used as tensors.
  def _tensor_conversion_aggregate(var, dtype=None, name=None, as_ref=False):
    return var._dense_var_to_tensor(dtype, name, as_ref)  # pylint: disable=protected-access
  
  
  ops.register_tensor_conversion_function(AggregatingVariable,
                                          _tensor_conversion_aggregate)
  ops.register_dense_tensor_like_type(AggregatingVariable)
  