code: # 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.
  # ==============================================================================
  """Class CollectiveAllReduceStrategy implementing DistributionStrategy."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import copy
  import weakref
  
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.core.protobuf import rewriter_config_pb2
  from tensorflow.core.protobuf import tensorflow_server_pb2
  from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib
  from tensorflow.python.distribute import cross_device_utils
  from tensorflow.python.distribute import device_util
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import input_lib
  from tensorflow.python.distribute import mirrored_strategy
  from tensorflow.python.distribute import multi_worker_util
  from tensorflow.python.distribute import numpy_dataset
  from tensorflow.python.distribute import reduce_util
  from tensorflow.python.distribute import values
  from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver
  from tensorflow.python.distribute.cluster_resolver import TFConfigClusterResolver
  from tensorflow.python.eager import context
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import collective_ops
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.util.tf_export import tf_export
  
  
  # TODO(yuefengz): support in-graph replication.
  @tf_export("distribute.experimental.MultiWorkerMirroredStrategy", v1=[])
  class CollectiveAllReduceStrategy(distribute_lib.Strategy):
    """A distribution strategy for synchronous training on multiple workers.
  
    This strategy implements synchronous distributed training across multiple
    workers, each with potentially multiple GPUs. Similar to
    `tf.distribute.MirroredStrategy`, it creates copies of all variables in the
    model on each device across all workers.
  
    It uses CollectiveOps's implementation of multi-worker all-reduce to
    to keep variables in sync. A collective op is a single op in the
    TensorFlow graph which can automatically choose an all-reduce algorithm in
    the TensorFlow runtime according to hardware, network topology and tensor
    sizes.
  
    By default it uses all local GPUs or CPU for single-worker training.
  
    When 'TF_CONFIG' environment variable is set, it parses cluster_spec,
    task_type and task_id from 'TF_CONFIG' and turns into a multi-worker strategy
    which mirrored models on GPUs of all machines in a cluster. In the current
    implementation, it uses all GPUs in a cluster and it assumes all workers have
    the same number of GPUs.
  
    You can also pass a `distribute.cluster_resolver.ClusterResolver` instance
    when instantiating the strategy. The task_type, task_id etc. will be parsed
    from the resolver instance instead of from the `TF_CONFIG` env var.
  
    It supports both eager mode and graph mode. However, for eager mode, it has to
    set up the eager context in its constructor and therefore all ops in eager
    mode have to run after the strategy object is created.
  
    """
    # TODO(anjalisridhar): Update our guides with examples showing how we can use
    # the cluster_resolver argument.
  
    def __init__(
        self,
        communication=cross_device_ops_lib.CollectiveCommunication.AUTO,
        cluster_resolver=None):
      """Creates the strategy.
  
      Args:
        communication: optional Enum of type
          `distribute.experimental.CollectiveCommunication`.  This provides a way
          for the user to override the choice of collective op communication.
          Possible values include `AUTO`, `RING`, and `NCCL`.
        cluster_resolver: optional `distribute.cluster_resolver.ClusterResolver`
          object. The default ClusterResolver that is used is the
          TFConfigClusterResolver which is instantiated from the TF_CONFIG env
          var.
      """
      # TODO(b/150151677): consider move communication to CollectiveHints.
      super(CollectiveAllReduceStrategy, self).__init__(
          CollectiveAllReduceExtended(
              self,
              communication=communication,
              cluster_resolver=cluster_resolver))
  
      distribute_lib.distribution_strategy_gauge.get_cell("V2").set(
          "MultiWorkerMirroredStrategy")
      # pylint: disable=protected-access
      distribute_lib.distribution_strategy_replica_gauge.get_cell(
          "num_workers").set(self.extended._num_workers)
      distribute_lib.distribution_strategy_replica_gauge.get_cell(
          "num_replicas_per_worker").set(self.extended._num_gpus_per_worker)
  
    @classmethod
    def _from_local_devices(
        cls,
        devices,
        communication=cross_device_ops_lib.CollectiveCommunication.AUTO):
      """A convenience method to create an object with a list of devices."""
      obj = cls(communication)
      obj.extended._initialize_local(TFConfigClusterResolver(), devices=devices)  # pylint: disable=protected-access
      return obj
  
    def scope(self):  # pylint: disable=useless-super-delegation
      """Returns a context manager selecting this Strategy as current.
  
      Inside a `with strategy.scope():` code block, this thread
      will use a variable creator set by `strategy`, and will
      enter its "cross-replica context".
  
      In `MultiWorkerMirroredStrategy`, all variables created inside
      `strategy.scope() will be mirrored on all replicas of each worker.
      Moreover, it also sets a default device scope so that ops without
      specified devices will end up on the correct worker.
  
      Returns:
        A context manager to use for creating variables with this strategy.
      """
      return super(CollectiveAllReduceStrategy, self).scope()
  
  
  @tf_export(v1=["distribute.experimental.MultiWorkerMirroredStrategy"])  # pylint: disable=missing-docstring
  class CollectiveAllReduceStrategyV1(distribute_lib.StrategyV1):
  
    __doc__ = CollectiveAllReduceStrategy.__doc__
  
    def __init__(
        self,
        communication=cross_device_ops_lib.CollectiveCommunication.AUTO,
        cluster_resolver=None):
      """Initializes the object."""
      super(CollectiveAllReduceStrategyV1, self).__init__(
          CollectiveAllReduceExtended(
              self,
              communication=communication,
              cluster_resolver=cluster_resolver))
      distribute_lib.distribution_strategy_gauge.get_cell("V1").set(
          "MultiWorkerMirroredStrategy")
      # pylint: disable=protected-access
      distribute_lib.distribution_strategy_replica_gauge.get_cell(
          "num_workers").set(self.extended._num_workers)
      distribute_lib.distribution_strategy_replica_gauge.get_cell(
          "num_gpu_per_worker").set(self.extended._num_gpus_per_worker)
  
  
  class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
    """Implementation of CollectiveAllReduceStrategy."""
  
    def __init__(self,
                 container_strategy,
                 communication,
                 cluster_resolver):
      self._cluster_resolver = cluster_resolver or TFConfigClusterResolver()
      distribute_lib.StrategyExtendedV1.__init__(self, container_strategy)
      assert isinstance(
          communication,
          cross_device_ops_lib.CollectiveCommunication)
      self._communication = communication
      self._initialize_strategy(self._cluster_resolver)
      self._cfer_fn_cache = weakref.WeakKeyDictionary()
      assert isinstance(self._cross_device_ops,
                        cross_device_ops_lib.CollectiveAllReduce)
  
    def _initialize_strategy(self, cluster_resolver):
      if cluster_resolver.cluster_spec().as_dict():
        self._initialize_multi_worker(cluster_resolver)
      else:
        self._initialize_local(cluster_resolver)
  
    def _initialize_local(self, cluster_resolver, devices=None):
      """Initializes the object for local training."""
      self._is_chief = True
      self._num_workers = 1
  
      if ops.executing_eagerly_outside_functions():
        try:
          context.context().configure_collective_ops(
              scoped_allocator_enabled_ops=("CollectiveReduce",))
        except RuntimeError:
          logging.warning("Collective ops is not configured at program startup. "
                          "Some performance features may not be enabled.")
        self._collective_ops_configured = True
  
      # TODO(b/126786766): TFConfigClusterResolver returns wrong number of GPUs in
      # some cases.
      if isinstance(cluster_resolver, TFConfigClusterResolver):
        num_gpus = context.num_gpus()
      else:
        num_gpus = cluster_resolver.num_accelerators().get("GPU", 0)
  
      if devices:
        local_devices = devices
      else:
        if num_gpus:
          local_devices = tuple("/device:GPU:%d" % i for i in range(num_gpus))
        else:
          local_devices = ("/device:CPU:0",)
  
      self._worker_device = device_util.canonicalize("/device:CPU:0")
      self._host_input_device = numpy_dataset.SingleDevice(self._worker_device)
  
      self._collective_keys = cross_device_utils.CollectiveKeys()
      self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
          devices=local_devices,
          group_size=len(local_devices),
          collective_keys=self._collective_keys,
          communication=self._communication)
      # CrossDeviceOps for per host tensors.
      self._host_cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
          devices=[self._worker_device],
          group_size=self._num_workers,
          collective_keys=self._collective_keys,
          communication=cross_device_ops_lib.CollectiveCommunication.RING,
      )
      super(CollectiveAllReduceExtended, self)._initialize_single_worker(
          local_devices)
  
      self._cluster_spec = None
      self._task_type = None
      self._task_id = None
  
      # This is a mark to tell whether we are running with standalone client or
      # independent worker. Right now with standalone client, strategy object is
      # created as local strategy and then turn into multi-worker strategy via
      # configure call.
      self._local_or_standalone_client_mode = True
  
      # Save the num_gpus_per_worker and rpc_layer for configure method.
      self._num_gpus_per_worker = num_gpus
      self._rpc_layer = cluster_resolver.rpc_layer
      self._warn_nccl_no_gpu()
  
      logging.info("Single-worker MultiWorkerMirroredStrategy with local_devices "
                   "= %r, communication = %s", local_devices, self._communication)
  
    def _initialize_multi_worker(self, cluster_resolver):
      """Initializes the object for multi-worker training."""
      cluster_spec = multi_worker_util.normalize_cluster_spec(
          cluster_resolver.cluster_spec())
      task_type = cluster_resolver.task_type
      task_id = cluster_resolver.task_id
      if task_type is None or task_id is None:
        raise ValueError("When `cluster_spec` is given, you must also specify "
                         "`task_type` and `task_id`.")
      self._cluster_spec = cluster_spec
      self._task_type = task_type
      self._task_id = task_id
  
      self._num_workers = multi_worker_util.worker_count(cluster_spec, task_type)
      if not self._num_workers:
        raise ValueError("No `worker`, `chief` or `evaluator` tasks can be found "
                         "in `cluster_spec`.")
  
      self._is_chief = multi_worker_util.is_chief(cluster_spec, task_type,
                                                  task_id)
  
      self._worker_device = "/job:%s/task:%d" % (task_type, task_id)
      self._host_input_device = numpy_dataset.SingleDevice(self._worker_device)
  
      if (ops.executing_eagerly_outside_functions() and
          not getattr(self, "_local_or_standalone_client_mode", False)):
        context.context().configure_collective_ops(
            collective_leader=multi_worker_util.collective_leader(
                cluster_spec, task_type, task_id),
            scoped_allocator_enabled_ops=("CollectiveReduce",),
            device_filters=("/job:%s/task:%d" % (task_type, task_id),))
        self._collective_ops_configured = True
  
      # Starting a std server in eager mode and in independent worker mode.
      if (context.executing_eagerly() and
          not getattr(self, "_std_server_started", False) and
          not getattr(self, "_local_or_standalone_client_mode", False)):
        # Checking _local_or_standalone_client_mode as well because we should not
        # create the std server in standalone client mode.
        config_proto = config_pb2.ConfigProto()
        config_proto = self._update_config_proto(config_proto)
  
        if hasattr(cluster_resolver, "port"):
          port = cluster_resolver.port
        else:
          port = 0
        server_def = tensorflow_server_pb2.ServerDef(
            cluster=cluster_spec.as_cluster_def(),
            default_session_config=config_proto,
            job_name=task_type,
            task_index=task_id,
            protocol=cluster_resolver.rpc_layer or "grpc",
            port=port)
        context.context().enable_collective_ops(server_def)
        self._std_server_started = True
        # The `ensure_initialized` is needed before calling
        # `context.context().devices()`.
        context.context().ensure_initialized()
        logging.info(
            "Enabled multi-worker collective ops with available devices: %r",
            context.context().devices())
  
      # TODO(yuefengz): The `num_gpus` is only for this particular task. It
      # assumes all workers have the same number of GPUs. We should remove this
      # assumption by querying all tasks for their numbers of GPUs.
      # TODO(b/126786766): TFConfigClusterResolver returns wrong number of GPUs in
      # some cases.
      if isinstance(cluster_resolver, TFConfigClusterResolver):
        num_gpus = context.num_gpus()
      else:
        num_gpus = cluster_resolver.num_accelerators().get("GPU", 0)
  
      if num_gpus:
        local_devices = tuple("%s/device:GPU:%d" % (self._worker_device, i)
                              for i in range(num_gpus))
      else:
        local_devices = (self._worker_device,)
  
      self._collective_keys = cross_device_utils.CollectiveKeys()
      self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
          devices=local_devices,
          group_size=len(local_devices) * self._num_workers,
          collective_keys=self._collective_keys,
          communication=self._communication)
      # CrossDeviceOps for per host tensors.
      self._host_cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
          devices=[self._worker_device],
          group_size=self._num_workers,
          collective_keys=self._collective_keys,
          communication=cross_device_ops_lib.CollectiveCommunication.RING,
      )
      super(CollectiveAllReduceExtended, self)._initialize_single_worker(
          local_devices)
      host_device = device_util.get_host_for_device(self._worker_device)
      self._input_workers = input_lib.InputWorkers(
          [(host_device, self.worker_devices)])
  
      # Add a default device so that ops without specified devices will not end up
      # on other workers.
      self._default_device = "/job:%s/task:%d" % (task_type, task_id)
  
      # Save the num_gpus_per_worker and rpc_layer for configure method.
      self._num_gpus_per_worker = num_gpus
      self._rpc_layer = cluster_resolver.rpc_layer
      self._warn_nccl_no_gpu()
  
      logging.info(
          "MultiWorkerMirroredStrategy with cluster_spec = %r, task_type = %r, "
          "task_id = %r, num_workers = %r, local_devices = %r, "
          "communication = %s", cluster_spec.as_dict(), task_type,
          task_id, self._num_workers, local_devices,
          self._communication)
  
    def _get_variable_creator_initial_value(self,
                                            replica_id,
                                            device,
                                            primary_var,
                                            **kwargs):
      if replica_id == 0:  # First replica on each worker.
        assert device is not None
        assert primary_var is None
  
        def initial_value_fn():  # pylint: disable=g-missing-docstring
          # Only the first device participates in the broadcast of initial values.
          group_key = self._collective_keys.get_group_key([device])
          group_size = self._num_workers
          collective_instance_key = (
              self._collective_keys.get_variable_instance_key())
  
          with ops.device(device):
            initial_value = kwargs["initial_value"]
            if callable(initial_value):
              initial_value = initial_value()
            assert not callable(initial_value)
            initial_value = ops.convert_to_tensor(
                initial_value, dtype=kwargs.get("dtype", None))
  
            if self._num_workers > 1:
              if self._is_chief:
                bcast_send = collective_ops.broadcast_send(
                    initial_value, initial_value.shape, initial_value.dtype,
                    group_size, group_key, collective_instance_key)
                with ops.control_dependencies([bcast_send]):
                  return array_ops.identity(initial_value)
              else:
                return collective_ops.broadcast_recv(initial_value.shape,
                                                     initial_value.dtype,
                                                     group_size, group_key,
                                                     collective_instance_key)
            return initial_value
  
        return initial_value_fn
      else:
        return super(CollectiveAllReduceExtended,
                     self)._get_variable_creator_initial_value(
                         replica_id=replica_id,
                         device=device,
                         primary_var=primary_var,
                         **kwargs)
  
    def _make_input_context(self):
      if self._cluster_spec is None:
        input_pipeline_id = 0
      else:
        input_pipeline_id = multi_worker_util.id_in_cluster(
            self._cluster_spec, self._task_type, self._task_id)
      input_context = distribute_lib.InputContext(
          num_input_pipelines=self._num_workers,
          input_pipeline_id=input_pipeline_id,
          num_replicas_in_sync=self._num_replicas_in_sync)
      return input_context
  
    def _experimental_distribute_dataset(self, dataset, options):
      input_context = self._make_input_context()
      return input_lib.get_distributed_dataset(
          dataset,
          self._input_workers,
          self._container_strategy(),
          split_batch_by=self._num_replicas_in_sync,
          input_context=input_context)
  
    def _experimental_distribute_datasets_from_function(self, dataset_fn,
                                                        options):
      input_context = self._make_input_context()
      return input_lib.get_distributed_datasets_from_function(
          dataset_fn=dataset_fn,
          input_workers=self._input_workers,
          input_contexts=[input_context],
          strategy=self._container_strategy())
  
    def _make_dataset_iterator(self, dataset):
      """Distributes the dataset to each local GPU."""
      input_context = self._make_input_context()
      return input_lib.DatasetIterator(
          dataset,
          self._input_workers,
          self._container_strategy(),
          split_batch_by=self._num_replicas_in_sync,
          input_context=input_context)
  
    def _make_input_fn_iterator(
        self,
        input_fn,
        replication_mode=distribute_lib.InputReplicationMode.PER_WORKER):
      """Distributes the input function to each local GPU."""
      input_context = self._make_input_context()
      return input_lib.InputFunctionIterator(input_fn, self._input_workers,
                                             [input_context],
                                             self._container_strategy())
  
    def _configure(self,
                   session_config=None,
                   cluster_spec=None,
                   task_type=None,
                   task_id=None):
      """Configures the object.
  
      Args:
        session_config: a `tf.compat.v1.ConfigProto`
        cluster_spec: a dict, ClusterDef or ClusterSpec object specifying the
          cluster configurations.
        task_type: the current task type, such as "worker".
        task_id: the current task id.
  
      Raises:
        ValueError: if `task_type` is not in the `cluster_spec`.
      """
      if cluster_spec:
        # Use the num_gpus_per_worker recorded in constructor since _configure
        # doesn't take num_gpus.
        cluster_resolver = SimpleClusterResolver(
            cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec),
            task_type=task_type,
            task_id=task_id,
            num_accelerators={"GPU": self._num_gpus_per_worker},
            rpc_layer=self._rpc_layer)
        self._initialize_multi_worker(cluster_resolver)
        assert isinstance(self._cross_device_ops,
                          cross_device_ops_lib.CollectiveAllReduce)
  
      if session_config:
        session_config.CopyFrom(self._update_config_proto(session_config))
  
    def _update_config_proto(self, config_proto):
      updated_config = copy.deepcopy(config_proto)
      # Enable the scoped allocator optimization for CollectiveOps.  This
      # optimization converts many small all-reduces into fewer larger
      # all-reduces.
      rewrite_options = updated_config.graph_options.rewrite_options
      rewrite_options.scoped_allocator_optimization = (
          rewriter_config_pb2.RewriterConfig.ON)
      # We turn on ScopedAllocator only for CollectiveReduce op, i.e. enable_op =
      # ["CollectiveReduce"].  Since we can't assign to a repeated proto field, we
      # clear and then append.
      del rewrite_options.scoped_allocator_opts.enable_op[:]
      rewrite_options.scoped_allocator_opts.enable_op.append("CollectiveReduce")
  
      if (not ops.executing_eagerly_outside_functions() and
          self._communication ==
          cross_device_ops_lib.CollectiveCommunication.NCCL):
        updated_config.experimental.collective_nccl = True
  
      if not self._cluster_spec:
        return updated_config
  
      assert self._task_type
      assert self._task_id is not None
  
      # Collective group leader is needed for collective ops to coordinate
      # workers.
      updated_config.experimental.collective_group_leader = (
          multi_worker_util.collective_leader(self._cluster_spec, self._task_type,
                                              self._task_id))
  
      # The device filters prevent communication between workers.
      del updated_config.device_filters[:]
      updated_config.device_filters.append(
          "/job:%s/task:%d" % (self._task_type, self._task_id))
  
      return updated_config
  
    def _get_cross_device_ops(self, value):
      # CollectiveAllReduce works on a predefined set of devices. In most cases
      # they should be the compute devices, but certain use cases may reduce host
      # tensors as well (e.g. early stopping). We infer the cross_device_ops to
      # use based on the number of devices, since inputs don't always have device
      # annotations. The compute devices one is preferred since we can potentially
      # leverage NCCL.
      if isinstance(value, values.DistributedValues):
        num_devices = len(value._values)  # pylint: disable=protected-access
      else:
        num_devices = 1
      if num_devices == len(self.worker_devices):
        return self._cross_device_ops
      else:
        return self._host_cross_device_ops
  
    def _reduce_to(self, reduce_op, value, destinations, experimental_hints):
      if (isinstance(value, values.Mirrored) and
          reduce_op == reduce_util.ReduceOp.MEAN):
        return value
      assert not isinstance(value, values.Mirrored)
  
      if (isinstance(value, values.DistributedValues) and
          len(self.worker_devices) == 1):
        value = value.values[0]
  
      # When there are multiple workers, we need to reduce across workers using
      # collective ops.
      if (not isinstance(value, values.DistributedValues) and
          self._num_workers == 1):
        # This function handles reducing values that are not PerReplica or
        # Mirrored values. For example, the same value could be present on all
        # replicas in which case `value` would be a single value or value could
        # be 0.
        return cross_device_ops_lib.reduce_non_distributed_value(
            reduce_op, value, destinations, len(self.worker_devices))
      return self._get_cross_device_ops(value).reduce(
          reduce_op,
          value,
          destinations=destinations,
          experimental_hints=experimental_hints)
  
    def _warn_nccl_no_gpu(self):
      if ((self._communication ==
           cross_device_ops_lib.CollectiveCommunication.NCCL) and
          self._num_gpus_per_worker == 0):
        logging.warning("Enabled NCCL communication but no GPUs detected/"
                        "specified.")
  
    def _in_multi_worker_mode(self):
      """Whether this strategy indicates working in multi-worker settings."""
      return self._num_workers > 1
  
    @property
    def experimental_between_graph(self):
      return True
  
    @property
    def experimental_should_init(self):
      return True
  
    @property
    def should_checkpoint(self):
      return self._is_chief
  
    @property
    def should_save_summary(self):
      return self._is_chief
  
    @property
    def _num_replicas_in_sync(self):
      return len(self.worker_devices) * self._num_workers
  
    # TODO(priyag): Delete this once all strategies use global batch size.
    @property
    def _global_batch_size(self):
      """`make_dataset_iterator` and `make_numpy_iterator` use global batch size.
  
      `make_input_fn_iterator` assumes per-replica batching.
  
      Returns:
        Boolean.
      """
      return True
   # Copyright 2015 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.
  # ==============================================================================
  
  """Operations for clipping (gradient, weight) tensors to min/max values."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import six
  
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import gen_nn_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.util import deprecation
  from tensorflow.python.util import dispatch
  from tensorflow.python.util.compat import collections_abc
  from tensorflow.python.util.tf_export import tf_export
  
  
  @tf_export("clip_by_value")
  @dispatch.add_dispatch_support
  def clip_by_value(t, clip_value_min, clip_value_max,
                    name=None):
    """Clips tensor values to a specified min and max.
  
    Given a tensor `t`, this operation returns a tensor of the same type and
    shape as `t` with its values clipped to `clip_value_min` and `clip_value_max`.
    Any values less than `clip_value_min` are set to `clip_value_min`. Any values
    greater than `clip_value_max` are set to `clip_value_max`.
  
    Note: `clip_value_min` needs to be smaller or equal to `clip_value_max` for
    correct results.
  
    For example:
  
    Basic usage passes a scalar as the min and max value.
  
    >>> t = tf.constant([[-10., -1., 0.], [0., 2., 10.]])
    >>> t2 = tf.clip_by_value(t, clip_value_min=-1, clip_value_max=1)
    >>> t2.numpy()
    array([[-1., -1.,  0.],
           [ 0.,  1.,  1.]], dtype=float32)
  
    The min and max can be the same size as `t`, or broadcastable to that size.
  
    >>> t = tf.constant([[-1, 0., 10.], [-1, 0, 10]])
    >>> clip_min = [[2],[1]]
    >>> t3 = tf.clip_by_value(t, clip_value_min=clip_min, clip_value_max=100)
    >>> t3.numpy()
    array([[ 2.,  2., 10.],
           [ 1.,  1., 10.]], dtype=float32)
  
    Broadcasting fails, intentionally, if you would expand the dimensions of `t`
  
    >>> t = tf.constant([[-1, 0., 10.], [-1, 0, 10]])
    >>> clip_min = [[[2, 1]]] # Has a third axis
    >>> t4 = tf.clip_by_value(t, clip_value_min=clip_min, clip_value_max=100)
    Traceback (most recent call last):
    ...
    InvalidArgumentError: Incompatible shapes: [2,3] vs. [1,1,2]
  
    It throws a `TypeError` if you try to clip an `int` to a `float` value
    (`tf.cast` the input to `float` first).
  
    >>> t = tf.constant([[1, 2], [3, 4]], dtype=tf.int32)
    >>> t5 = tf.clip_by_value(t, clip_value_min=-3.1, clip_value_max=3.1)
    Traceback (most recent call last):
    ...
    TypeError: Cannot convert ...
  
  
    Args:
      t: A `Tensor` or `IndexedSlices`.
      clip_value_min: The minimum value to clip to. A scalar `Tensor` or one that
        is broadcastable to the shape of `t`.
      clip_value_max: The minimum value to clip to. A scalar `Tensor` or one that
        is broadcastable to the shape of `t`.
      name: A name for the operation (optional).
  
    Returns:
      A clipped `Tensor` or `IndexedSlices`.
  
    Raises:
      `tf.errors.InvalidArgumentError`: If the clip tensors would trigger array
        broadcasting that would make the returned tensor larger than the input.
      TypeError: If dtype of the input is `int32` and dtype of
        the `clip_value_min` or `clip_value_max` is `float32`
    """
    with ops.name_scope(name, "clip_by_value",
                        [t, clip_value_min, clip_value_max]) as name:
      values = ops.convert_to_tensor(
          t.values if isinstance(t, ops.IndexedSlices) else t, name="t")
  
      # Go through list of tensors, for each value in each tensor clip
      t_min = math_ops.minimum(values, clip_value_max)
      # Assert that the shape is compatible with the initial shape,
      # to prevent unintentional broadcasting.
      _ = values.shape.merge_with(t_min.shape)
  
      t_max = math_ops.maximum(t_min, clip_value_min, name=name)
      _ = values.shape.merge_with(t_max.shape)
  
      if isinstance(t, ops.IndexedSlices):
        t_max = ops.IndexedSlices(t_max, t.indices, t.dense_shape)
  
    return t_max
    # TODO(scottzhu): switch to use new implementation in 2 weeks.
    # return gen_math_ops.clip_by_value(
    #     t, clip_value_min, clip_value_max, name=name)
  
  
  # TODO(scottzhu): switch to use new implementation in 2 weeks.
  # @ops.RegisterGradient("ClipByValue")
  def _clip_by_value_grad(op, grad):
    """Returns grad of clip_by_value."""
    x = op.inputs[0]
    y = op.inputs[1]
    z = op.inputs[2]
    gdtype = grad.dtype
    sx = array_ops.shape(x)
    sy = array_ops.shape(y)
    sz = array_ops.shape(z)
    gradshape = array_ops.shape(grad)
    zeros = array_ops.zeros(gradshape, gdtype)
    xymask = math_ops.less(x, y)
    xzmask = math_ops.greater(x, z)
    rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy)
    rx, rz = gen_array_ops.broadcast_gradient_args(sx, sz)
    xgrad = array_ops.where(math_ops.logical_or(xymask, xzmask), zeros, grad)
    ygrad = array_ops.where(xymask, grad, zeros)
    zgrad = array_ops.where(xzmask, grad, zeros)
    gx = array_ops.reshape(math_ops.reduce_sum(xgrad, rx), sx)
    gy = array_ops.reshape(math_ops.reduce_sum(ygrad, ry), sy)
    gz = array_ops.reshape(math_ops.reduce_sum(zgrad, rz), sz)
    return (gx, gy, gz)
  
  
  @tf_export("clip_by_norm")
  @dispatch.add_dispatch_support
  def clip_by_norm(t, clip_norm, axes=None, name=None):
    """Clips tensor values to a maximum L2-norm.
  
    Given a tensor `t`, and a maximum clip value `clip_norm`, this operation
    normalizes `t` so that its L2-norm is less than or equal to `clip_norm`,
    along the dimensions given in `axes`. Specifically, in the default case
    where all dimensions are used for calculation, if the L2-norm of `t` is
    already less than or equal to `clip_norm`, then `t` is not modified. If
    the L2-norm is greater than `clip_norm`, then this operation returns a
    tensor of the same type and shape as `t` with its values set to:
  
    `t * clip_norm / l2norm(t)`
  
    In this case, the L2-norm of the output tensor is `clip_norm`.
  
    As another example, if `t` is a matrix and `axes == [1]`, then each row
    of the output will have L2-norm less than or equal to `clip_norm`. If
    `axes == [0]` instead, each column of the output will be clipped.
  
    Code example:
  
    >>> some_nums = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.float32)
    >>> tf.clip_by_norm(some_nums, 2.0).numpy()
    array([[0.26967996, 0.5393599 , 0.80903983, 1.0787199 , 1.3483998 ]],
          dtype=float32)
  
    This operation is typically used to clip gradients before applying them with
    an optimizer.  Most gradient data is a collection of different shaped tensors
    for different parts of the model.  Thus, this is a common usage:
  
    ```
    # Get your gradients after training
    loss_value, grads = grad(model, features, labels)
  
    # Apply some clipping
    grads = [tf.clip_by_norm(g, norm)
                 for g in grads]
  
    # Continue on with training
    optimizer.apply_gradients(grads)
    ```
  
    Args:
      t: A `Tensor` or `IndexedSlices`.  This must be a floating point type.
      clip_norm: A 0-D (scalar) `Tensor` > 0. A maximum clipping value, also
        floating point
      axes: A 1-D (vector) `Tensor` of type int32 containing the dimensions
        to use for computing the L2-norm. If `None` (the default), uses all
        dimensions.
      name: A name for the operation (optional).
  
    Returns:
      A clipped `Tensor` or `IndexedSlices`.
  
    Raises:
      ValueError: If the clip_norm tensor is not a 0-D scalar tensor.
      TypeError: If dtype of the input is not a floating point or
        complex type.
    """
    with ops.name_scope(name, "clip_by_norm", [t, clip_norm]) as name:
      values = ops.convert_to_tensor(
          t.values if isinstance(t, ops.IndexedSlices) else t, name="t")
  
      # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
      l2sum = math_ops.reduce_sum(values * values, axes, keepdims=True)
      pred = l2sum > 0
      # Two-tap tf.where trick to bypass NaN gradients
      l2sum_safe = array_ops.where(pred, l2sum, array_ops.ones_like(l2sum))
      l2norm = array_ops.where(pred, math_ops.sqrt(l2sum_safe), l2sum)
      intermediate = values * clip_norm
      # Assert that the shape is compatible with the initial shape,
      # to prevent unintentional broadcasting.
      _ = values.shape.merge_with(intermediate.shape)
      values_clip = array_ops.identity(
          intermediate / math_ops.maximum(l2norm, clip_norm), name=name)
  
      if isinstance(t, ops.IndexedSlices):
        return ops.IndexedSlices(values_clip, t.indices, t.dense_shape)
  
      return values_clip
  
  
  @tf_export("linalg.global_norm", v1=["linalg.global_norm", "global_norm"])
  @dispatch.add_dispatch_support
  @deprecation.deprecated_endpoints("global_norm")
  def global_norm(t_list, name=None):
    """Computes the global norm of multiple tensors.
  
    Given a tuple or list of tensors `t_list`, this operation returns the
    global norm of the elements in all tensors in `t_list`. The global norm is
    computed as:
  
    `global_norm = sqrt(sum([l2norm(t)**2 for t in t_list]))`
  
    Any entries in `t_list` that are of type None are ignored.
  
    Args:
      t_list: A tuple or list of mixed `Tensors`, `IndexedSlices`, or None.
      name: A name for the operation (optional).
  
    Returns:
      A 0-D (scalar) `Tensor` of type `float`.
  
    Raises:
      TypeError: If `t_list` is not a sequence.
    """
    if (not isinstance(t_list, collections_abc.Sequence) or
        isinstance(t_list, six.string_types)):
      raise TypeError("t_list should be a sequence")
    t_list = list(t_list)
    with ops.name_scope(name, "global_norm", t_list) as name:
      values = [
          ops.convert_to_tensor(
              t.values if isinstance(t, ops.IndexedSlices) else t,
              name="t_%d" % i)
          if t is not None else t
          for i, t in enumerate(t_list)]
      half_squared_norms = []
      for v in values:
        if v is not None:
          with ops.colocate_with(v):
            half_squared_norms.append(gen_nn_ops.l2_loss(v))
  
      half_squared_norm = math_ops.reduce_sum(array_ops.stack(half_squared_norms))
  
      norm = math_ops.sqrt(
          half_squared_norm *
          constant_op.constant(2.0, dtype=half_squared_norm.dtype),
          name="global_norm")
  
    return norm
  
  
  @tf_export("clip_by_global_norm")
  @dispatch.add_dispatch_support
  def clip_by_global_norm(t_list, clip_norm, use_norm=None, name=None):
    """Clips values of multiple tensors by the ratio of the sum of their norms.
  
    Given a tuple or list of tensors `t_list`, and a clipping ratio `clip_norm`,
    this operation returns a list of clipped tensors `list_clipped`
    and the global norm (`global_norm`) of all tensors in `t_list`. Optionally,
    if you've already computed the global norm for `t_list`, you can specify
    the global norm with `use_norm`.
  
    To perform the clipping, the values `t_list[i]` are set to:
  
        t_list[i] * clip_norm / max(global_norm, clip_norm)
  
    where:
  
        global_norm = sqrt(sum([l2norm(t)**2 for t in t_list]))
  
    If `clip_norm > global_norm` then the entries in `t_list` remain as they are,
    otherwise they're all shrunk by the global ratio.
  
    If `global_norm == infinity` then the entries in `t_list` are all set to `NaN`
    to signal that an error occurred.
  
    Any of the entries of `t_list` that are of type `None` are ignored.
  
    This is the correct way to perform gradient clipping (Pascanu et al., 2012).
  
    However, it is slower than `clip_by_norm()` because all the parameters must be
    ready before the clipping operation can be performed.
  
    Args:
      t_list: A tuple or list of mixed `Tensors`, `IndexedSlices`, or None.
      clip_norm: A 0-D (scalar) `Tensor` > 0. The clipping ratio.
      use_norm: A 0-D (scalar) `Tensor` of type `float` (optional). The global
        norm to use. If not provided, `global_norm()` is used to compute the norm.
      name: A name for the operation (optional).
  
    Returns:
      list_clipped: A list of `Tensors` of the same type as `list_t`.
      global_norm: A 0-D (scalar) `Tensor` representing the global norm.
  
    Raises:
      TypeError: If `t_list` is not a sequence.
  
    References:
      On the difficulty of training Recurrent Neural Networks:
        [Pascanu et al., 2012](http://proceedings.mlr.press/v28/pascanu13.html)
        ([pdf](http://proceedings.mlr.press/v28/pascanu13.pdf))
    """
    if (not isinstance(t_list, collections_abc.Sequence) or
        isinstance(t_list, six.string_types)):
      raise TypeError("t_list should be a sequence")
    t_list = list(t_list)
    if use_norm is None:
      use_norm = global_norm(t_list, name)
  
    with ops.name_scope(name, "clip_by_global_norm",
                        t_list + [clip_norm]) as name:
      # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
      scale_for_finite = clip_norm * math_ops.minimum(
          1.0 / use_norm,
          constant_op.constant(1.0, dtype=use_norm.dtype) / clip_norm)
      # If use_norm is any finite number, this is a no-op. For inf/-inf/NaN,
      # this will make scale NaN.
      scale = scale_for_finite + (use_norm - use_norm)
  
      values = [
          ops.convert_to_tensor(
              t.values if isinstance(t, ops.IndexedSlices) else t,
              name="t_%d" % i)
          if t is not None else t
          for i, t in enumerate(t_list)]
  
      values_clipped = []
      for i, v in enumerate(values):
        if v is None:
          values_clipped.append(None)
        else:
          with ops.colocate_with(v):
            values_clipped.append(
                array_ops.identity(v * scale, name="%s_%d" % (name, i)))
  
      list_clipped = [
          ops.IndexedSlices(c_v, t.indices, t.dense_shape)
          if isinstance(t, ops.IndexedSlices)
          else c_v
          for (c_v, t) in zip(values_clipped, t_list)]
  
    return list_clipped, use_norm
  
  
  @deprecation.deprecated(
      date=None,
      instructions="clip_by_average_norm is deprecated in TensorFlow 2.0. Please "
      "use clip_by_norm(t, clip_norm * tf.cast(tf.size(t), tf.float32), name) "
      "instead.")
  @tf_export(v1=["clip_by_average_norm"])
  @dispatch.add_dispatch_support
  def clip_by_average_norm(t, clip_norm, name=None):
    """Clips tensor values to a maximum average L2-norm.
  
    Given a tensor `t`, and a maximum clip value `clip_norm`, this operation
    normalizes `t` so that its average L2-norm is less than or equal to
    `clip_norm`. Specifically, if the average L2-norm is already less than or
    equal to `clip_norm`, then `t` is not modified. If the average L2-norm is
    greater than `clip_norm`, then this operation returns a tensor of the same
    type and shape as `t` with its values set to:
  
    `t * clip_norm / l2norm_avg(t)`
  
    In this case, the average L2-norm of the output tensor is `clip_norm`.
  
    This operation is typically used to clip gradients before applying them with
    an optimizer.
  
    Args:
      t: A `Tensor`.
      clip_norm: A 0-D (scalar) `Tensor` > 0. A maximum clipping value.
      name: A name for the operation (optional).
  
    Returns:
      A clipped `Tensor`.
    """
    with ops.name_scope(name, "clip_by_average_norm", [t, clip_norm]) as name:
      t = ops.convert_to_tensor(t, name="t")
  
      # Calculate L2-norm per element, clip elements by ratio of clip_norm to
      # L2-norm per element
      n_element = math_ops.cast(array_ops.size(t), dtypes.float32)
      l2norm_inv = math_ops.rsqrt(
          math_ops.reduce_sum(t * t, math_ops.range(array_ops.rank(t))))
      tclip = array_ops.identity(
          t * clip_norm * math_ops.minimum(
              l2norm_inv * n_element, constant_op.constant(1.0) / clip_norm),
          name=name)
  
    return tclip
  
apis: tensorflow.python.distribute.multi_worker_util.collective_leader tensorflow.python.distribute.multi_worker_util.worker_count tensorflow.python.distribute.multi_worker_util.is_chief tensorflow.python.distribute.multi_worker_util.id_in_cluster tensorflow.python.distribute.distribute_lib.StrategyExtendedV1.__init__ tensorflow.python.distribute.input_lib.InputWorkers tensorflow.python.distribute.distribute_lib.InputContext tensorflow.python.ops.collective_ops.broadcast_recv tensorflow.python.eager.context.num_gpus tensorflow.python.distribute.distribute_lib.distribution_strategy_replica_gauge.get_cell tensorflow.python.distribute.cross_device_utils.CollectiveKeys tensorflow.python.ops.collective_ops.broadcast_send tensorflow.python.util.tf_export.tf_export tensorflow.python.distribute.device_util.canonicalize tensorflow.python.platform.tf_logging.warning tensorflow.python.framework.ops.device tensorflow.python.distribute.multi_worker_util.normalize_cluster_spec tensorflow.python.platform.tf_logging.info tensorflow.python.ops.array_ops.identity tensorflow.python.eager.context.context tensorflow.python.framework.ops.control_dependencies tensorflow.python.distribute.numpy_dataset.SingleDevice tensorflow.python.distribute.cross_device_ops.CollectiveAllReduce tensorflow.python.distribute.distribute_lib.distribution_strategy_gauge.get_cell tensorflow.core.protobuf.config_pb2.ConfigProto tensorflow.python.framework.ops.executing_eagerly_outside_functions tensorflow.python.distribute.device_util.get_host_for_device tensorflow.python.eager.context.executing_eagerly tensorflow.python.distribute.cluster_resolver.TFConfigClusterResolver
apis: tensorflow.python.util.deprecation.deprecated tensorflow.python.framework.ops.colocate_with tensorflow.python.ops.math_ops.maximum tensorflow.python.util.deprecation.deprecated_endpoints tensorflow.python.framework.constant_op.constant tensorflow.python.framework.ops.IndexedSlices tensorflow.python.util.tf_export.tf_export tensorflow.python.ops.gen_nn_ops.l2_loss tensorflow.python.ops.math_ops.greater tensorflow.python.ops.array_ops.rank tensorflow.python.ops.math_ops.minimum tensorflow.python.ops.math_ops.less tensorflow.python.ops.array_ops.zeros tensorflow.python.ops.array_ops.shape tensorflow.python.ops.math_ops.reduce_sum tensorflow.python.ops.array_ops.identity tensorflow.python.ops.math_ops.logical_or tensorflow.python.framework.ops.convert_to_tensor tensorflow.python.ops.math_ops.sqrt tensorflow.python.framework.ops.name_scope tensorflow.python.ops.array_ops.ones_like tensorflow.python.ops.array_ops.stack tensorflow.python.ops.array_ops.where tensorflow.python.ops.array_ops.size tensorflow.python.ops.gen_array_ops.broadcast_gradient_args