chore: import upstream snapshot with attribution

tensorflow/python/training/experimental/BUILD

@@ -0,0 +1,148 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "loss_scale",
+    srcs = ["loss_scale.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "loss_scale_optimizer",
+    srcs = ["loss_scale_optimizer.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:smart_cond",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/training:optimizer",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "loss_scale_optimizer_test",
+    size = "small",
+    srcs = ["loss_scale_optimizer_test.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale",
+        ":loss_scale_optimizer",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:mirrored_strategy",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:custom_gradient",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:gradient_descent",
+        "//tensorflow/python/training:momentum",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_test(
+    name = "loss_scale_test",
+    size = "medium",
+    srcs = ["loss_scale_test.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:mirrored_strategy",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "mixed_precision_global_state",
+    srcs = ["mixed_precision_global_state.py"],
+    strict_deps = True,
+    deps = ["//tensorflow/python/util:tf_export"],
+)
+
+py_library(
+    name = "mixed_precision",
+    srcs = ["mixed_precision.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale_optimizer",
+        ":mixed_precision_global_state",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:optimizer",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "mixed_precision_test",
+    size = "small",
+    srcs = ["mixed_precision_test.py"],
+    deps = [
+        ":loss_scale_optimizer",
+        ":mixed_precision",
+        ":mixed_precision_global_state",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:gradient_descent",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/training/experimental/loss_scale.py

@@ -0,0 +1,453 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains LossScale classes."""
+import abc
+
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@deprecation.deprecated_endpoints('mixed_precision.experimental.LossScale',
+                                  'train.experimental.LossScale')
+@tf_export(
+    v1=[
+        'mixed_precision.LossScale',
+        'mixed_precision.experimental.LossScale',
+        'train.experimental.LossScale'
+    ])
+class LossScale(trackable.Trackable, metaclass=abc.ABCMeta):
+  """Base class for all TF1 loss scales.
+
+  This is an abstract base class, so you cannot instantiate it directly.
+  Instead, use one of its concrete subclasses:
+    * `tf.compat.v1.mixed_precision.DynamicLossScale`
+    * `tf.compat.v1.mixed_precision.FixedLossScale`
+
+  Loss scaling is a process that multiplies the loss by a multiplier called the
+  loss scale, and divides each gradient by the same multiplier. The pseudocode
+  for this process is:
+
+  ```
+  loss = ...
+  loss *= loss_scale
+  grads = gradients(loss, vars)
+  grads /= loss_scale
+  ```
+
+  Mathematically, loss scaling has no effect, but can help avoid numerical
+  underflow in intermediate gradients when float16 tensors are used for mixed
+  precision training. By multiplying the loss, each intermediate gradient will
+  have the same multiplier applied.
+
+  Instances of this class represent a loss scale. Calling instances of this
+  class returns the loss scale as a scalar float32 tensor, while method
+  `update()` updates the loss scale depending on the values of the gradients.
+  Optimizers use instances of this class to scale loss and gradients.
+
+  In most functions that accept a LossScale, you can also pass an int (such as
+  8) to create a `FixedLossScale` or the string `"dynamic"` to create a dynamic
+  loss scale.
+  """
+
+  def __init__(self):
+    """Initializes the loss scale class."""
+    self._weights = {}
+
+  @abc.abstractmethod
+  def __call__(self):
+    """Returns the current loss scale as a scalar `float32` tensor."""
+    pass
+
+  @abc.abstractmethod
+  def update(self, grads):
+    """Updates the value of the loss scale.
+
+    The loss scale will be potentially updated, based on the value of `grads`.
+    The tensor returned by calling this class is only updated when this function
+    is evaluated.
+
+    In eager mode, this directly updates the loss scale, so that calling
+    `__call__` will return the newly updated loss scale. In graph mode,
+    this returns an op that, when evaluated, updates the loss scale.
+
+    This function also returns a `should_apply_gradients` bool. If False,
+    gradients should not be applied to the variables that step, as nonfinite
+    gradients were found, and the loss scale has been be updated to reduce the
+    chance of finding nonfinite gradients in the next step. Some loss scale
+    classes will always return True, as they cannot adjust themselves in
+    response to nonfinite gradients.
+
+    When a DistributionStrategy is used, this function may only be called in a
+    cross-replica context.
+
+    Args:
+      grads: A nested structure of unscaled gradients, each which is the
+        gradient of the loss with respect to a weight. The gradients should have
+        already been divided by the loss scale being before passed to this
+        function. 'None' gradients are accepted, and are ignored.
+
+    Returns:
+      update_op: In eager mode, None. In graph mode, an op to update the loss
+        scale.
+      should_apply_gradients: Either a bool or a scalar boolean tensor. If
+        False, the caller should skip applying `grads` to the variables this
+        step.
+    """
+    pass
+
+  def _add_weight(self, name, initial_value, dtype=None):
+    """Adds a weight to this loss scale.
+
+    Args:
+      name: Variable name.
+      initial_value: The variable's initial value.
+      dtype: The type of the variable.
+
+    Returns:
+      A variable.
+
+    Raises:
+      RuntimeError: If a weight with `name` has already been added.
+    """
+    variable = variable_v1.VariableV1(
+        initial_value=initial_value,
+        name=name,
+        dtype=dtype,
+        trainable=False,
+        use_resource=True,
+        synchronization=variables.VariableSynchronization.AUTO,
+        # Set aggregation to NONE, as loss scaling variables should never be
+        # aggregated.
+        aggregation=variables.VariableAggregation.NONE)
+    if context.executing_eagerly():
+      graph_key = None
+    else:
+      graph = ops.get_default_graph()
+      graph_key = graph._graph_key  # pylint: disable=protected-access
+
+    key = (name, graph_key)
+    if self._weights.get(key, None) is not None:
+      raise RuntimeError('Duplicate variables detected. {}'.format(key))
+    self._weights[key] = variable
+    self._handle_deferred_dependencies(name=name, trackable=variable)
+    return variable
+
+  def _trackable_children(self,
+                          save_type=trackable.SaveType.CHECKPOINT,
+                          **kwargs):
+    """From Trackable. Gather graph-specific weights to save."""
+    if context.executing_eagerly():
+      graph_key = None
+    else:
+      graph = ops.get_default_graph()
+      graph_key = graph._graph_key  # pylint: disable=protected-access
+    weights = {}
+    for (name, g), v in sorted(self._weights.items(), key=lambda i: i[0][0]):
+      if g == graph_key:
+        weights[name] = v
+    weights.update(
+        super(LossScale, self)._trackable_children(save_type, **kwargs))
+    return weights
+
+  def _lookup_dependency(self, name, cached_dependencies=None):
+    """From Trackable. Find a weight in the current graph."""
+    unconditional = super(LossScale, self)._lookup_dependency(
+        name, cached_dependencies)
+    if unconditional is not None:
+      return unconditional
+    if context.executing_eagerly():
+      graph_key = None
+    else:
+      graph = ops.get_default_graph()
+      graph_key = graph._graph_key  # pylint: disable=protected-access
+    return self._weights.get((name, graph_key), None)
+
+  @abc.abstractmethod
+  def get_config(self):
+    """Returns the config of this loss scale."""
+    pass
+
+  @classmethod
+  def from_config(cls, config):
+    """Creates the LossScale from its config."""
+    return cls(**config)
+
+
+@deprecation.deprecated_endpoints('mixed_precision.experimental.FixedLossScale',
+                                  'train.experimental.FixedLossScale')
+@tf_export(
+    v1=[
+        'mixed_precision.FixedLossScale',
+        'mixed_precision.experimental.FixedLossScale',
+        'train.experimental.FixedLossScale'
+    ])
+class FixedLossScale(LossScale):
+  """Loss scale with a fixed value.
+
+  The loss scale is not updated for the lifetime of instances of this class.
+  A given instance of this class always returns the same number when called.
+  """
+
+  @deprecation.deprecated(
+      None, 'Use tf.keras.mixed_precision.LossScaleOptimizer instead. '
+            'LossScaleOptimizer now has all the functionality of '
+            'FixedLossScale')
+  def __init__(self, loss_scale_value):
+    """Creates the fixed loss scale.
+
+    Args:
+      loss_scale_value: A Python float. Its ideal value varies depending on
+        models to run. Choosing a too small loss_scale might affect model
+        quality; a too big loss_scale might cause inf or nan. There is no single
+        right loss_scale to apply. There is no harm choosing a relatively big
+        number as long as no nan or inf is encountered in training.
+
+    Raises:
+      ValueError: If loss_scale_value is less than 1.
+    """
+    super(FixedLossScale, self).__init__()
+    if not isinstance(loss_scale_value, (int, float)):
+      raise ValueError('loss_scale_value must be a Python int or float.')
+    if loss_scale_value < 1:
+      raise ValueError('loss_scale_value must be at least 1.')
+    # It's important we do not create tensors in the constructor, as such
+    # tensors might be on a different device or tf.function vs when the tensor
+    # is used. This would hurt performance. Therefore, we do not create a tensor
+    # from loss_scale_value, but instead leave it as a Python float.
+    # TODO(reedwm): Also do not create tensors in the DynamicLossScale
+    # constructor.
+    self._loss_scale_value = float(loss_scale_value)
+
+  def __call__(self):
+    return ops.convert_to_tensor(self._loss_scale_value)
+
+  def update(self, grads):
+    del grads
+    return control_flow_ops.no_op(), True
+
+  def __repr__(self):
+    return 'FixedLossScale(%s)' % self._loss_scale_value
+
+  def get_config(self):
+    return {'loss_scale_value': self._loss_scale_value}
+
+
+def _is_all_finite(grads):
+  """Returns a scalar boolean tensor indicating if all gradients are finite."""
+  def raw_values(g):
+    return g.values if isinstance(g, indexed_slices.IndexedSlices) else g
+
+  is_finite_per_grad = [
+      math_ops.reduce_all(math_ops.is_finite(raw_values(g)))
+      for g in grads
+      if g is not None
+  ]
+  return math_ops.reduce_all(is_finite_per_grad)
+
+
+def _op_in_graph_mode(tensor):
+  """Returns the tensor's op in graph mode, or the tensor in eager mode.
+
+  This is useful because sometimes an op is needed in graph mode instead of a
+  tensor. In eager mode, there are no ops.
+
+  Args:
+    tensor: A tensor.
+
+  Returns:
+    The tensor's op in graph mode. The tensor in eager mode.
+  """
+  if context.executing_eagerly():
+    return tensor
+  return tensor.op
+
+
+def _assign_if_finite(var, value):
+  """Assigns a value to a variable if the value is finite."""
+  return cond.cond(
+      math_ops.is_finite(value), lambda: _op_in_graph_mode(var.assign(value)),
+      control_flow_ops.no_op)
+
+
+@deprecation.deprecated_endpoints(
+    'mixed_precision.experimental.DynamicLossScale',
+    'train.experimental.DynamicLossScale')
+@tf_export(
+    v1=[
+        'mixed_precision.DynamicLossScale',
+        'mixed_precision.experimental.DynamicLossScale',
+        'train.experimental.DynamicLossScale'
+    ])
+class DynamicLossScale(LossScale):
+  """Loss scale that dynamically adjusts itself.
+
+  Dynamic loss scaling works by adjusting the loss scale as training progresses.
+  The goal is to keep the loss scale as high as possible without overflowing the
+  gradients. As long as the gradients do not overflow, raising the loss scale
+  never hurts.
+
+  The algorithm starts by setting the loss scale to an initial value. Every N
+  steps that the gradients are finite, the loss scale is increased by some
+  factor. However, if a NaN or Inf gradient is found, the gradients for that
+  step are not applied, and the loss scale is decreased by the factor. This
+  process tends to keep the loss scale as high as possible without gradients
+  overflowing.
+  """
+
+  @deprecation.deprecated(
+      None, 'Use tf.keras.mixed_precision.LossScaleOptimizer instead. '
+            'LossScaleOptimizer now has all the functionality of '
+            'DynamicLossScale')
+  def __init__(self,
+               initial_loss_scale=2 ** 15,  # See docstring for why this is big.
+               increment_period=2000,
+               multiplier=2.):
+    """Creates the dynamic loss scale.
+
+    Args:
+      initial_loss_scale: A Python float.  The loss scale to use at the
+        beginning. It's better to start this at a very high number, because a
+        loss scale that is too high gets lowered far more quickly than a loss
+        scale that is too low gets raised. The default is 2 ** 15, which is
+        approximately half the maximum float16 value.
+      increment_period: Increases loss scale every `increment_period`
+        consecutive steps that finite gradients are encountered. If a nonfinite
+        gradient is encountered, the count is reset back to zero.
+      multiplier: The multiplier to use when increasing or decreasing the loss
+        scale.
+    """
+    super(DynamicLossScale, self).__init__()
+    self._initial_loss_scale = float(initial_loss_scale)
+    self._increment_period = int(increment_period)
+    self._multiplier = float(multiplier)
+
+    self._current_loss_scale = self._add_weight(
+        name='current_loss_scale',
+        dtype=dtypes.float32,
+        initial_value=self._initial_loss_scale)
+    # The number of consecutive steps with finite gradients since the last
+    # nonfinite gradient or change in loss scale.
+    self._num_good_steps = self._add_weight(
+        name='good_steps', dtype=dtypes.int64, initial_value=0)
+
+  @property
+  def initial_loss_scale(self):
+    return self._initial_loss_scale
+
+  @property
+  def increment_period(self):
+    return self._increment_period
+
+  @property
+  def multiplier(self):
+    return self._multiplier
+
+  def __call__(self):
+    return ops.convert_to_tensor(self._current_loss_scale)
+
+  def update(self, grads):
+    """Updates loss scale based on if gradients are finite in current step."""
+    grads = nest.flatten(grads)
+    if distribute_lib.has_strategy():
+      distribution = distribute_lib.get_cross_replica_context()
+
+      def get_is_finite(grads):
+        is_finite = _is_all_finite(grads)
+        # We cast to float, because we cannot reduce booleans with
+        # DistributionStrategy.
+        return math_ops.cast(is_finite, dtypes.float32)
+
+      is_finite_float = distribution.extended.call_for_each_replica(
+          get_is_finite, args=(grads,))
+      reduced_is_finite_float = distribution.reduce(reduce_util.ReduceOp.SUM,
+                                                    is_finite_float, axis=None)
+      is_finite = math_ops.equal(reduced_is_finite_float,
+                                 distribution.num_replicas_in_sync)
+    else:
+      is_finite = _is_all_finite(grads)
+
+    def update_if_finite_grads():
+      """Update assuming the gradients are finite."""
+
+      def incr_loss_scale():
+        new_loss_scale = self._current_loss_scale * self._multiplier
+        return control_flow_ops.group(
+            _assign_if_finite(self._current_loss_scale, new_loss_scale),
+            self._num_good_steps.assign(0))
+
+      return cond.cond(
+          self._num_good_steps + 1 >= self._increment_period,
+          incr_loss_scale, lambda: _op_in_graph_mode(
+              self._num_good_steps.assign_add(1)))
+
+    def update_if_not_finite_grads():
+      """Update assuming the gradients are nonfinite."""
+
+      new_loss_scale = math_ops.maximum(
+          self._current_loss_scale / self._multiplier, 1)
+      return control_flow_ops.group(
+          self._num_good_steps.assign(0),
+          self._current_loss_scale.assign(new_loss_scale))
+
+    update_op = cond.cond(is_finite, update_if_finite_grads,
+                          update_if_not_finite_grads)
+    should_apply_gradients = is_finite
+    return update_op, should_apply_gradients
+
+  def __repr__(self):
+    if context.executing_eagerly():
+      return ('DynamicLossScale(current_loss_scale=%s, num_good_steps=%s, '
+              'initial_loss_scale=%s, increment_period=%s, multiplier=%s)' %
+              (self._current_loss_scale.numpy(), self._num_good_steps.numpy(),
+               self.initial_loss_scale, self.increment_period, self.multiplier))
+    else:
+      return ('DynamicLossScale(initial_loss_scale=%s, increment_period=%s, '
+              'multiplier=%s)' %
+              (self.initial_loss_scale, self.increment_period, self.multiplier))
+
+  def get_config(self):
+    return {
+        'initial_loss_scale': self.initial_loss_scale,
+        'increment_period': self.increment_period,
+        'multiplier': self.multiplier,
+    }
+
+
+def get(identifier):
+  """Get a loss scale object."""
+  if isinstance(identifier, (int, float)):
+    return FixedLossScale(identifier)
+  if identifier == 'dynamic':
+    return DynamicLossScale()
+  if isinstance(identifier, LossScale):
+    return identifier
+  elif identifier is None:
+    return None
+  else:
+    raise ValueError('Could not interpret loss scale identifier: %s' %
+                     identifier)

tensorflow/python/training/experimental/loss_scale_optimizer.py

@@ -0,0 +1,251 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains LossScale classes."""
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import smart_cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.training.experimental import loss_scale as loss_scale_module
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@deprecation.deprecated_endpoints(
+    'train.experimental.MixedPrecisionLossScaleOptimizer')
+@tf_export(v1=['mixed_precision.MixedPrecisionLossScaleOptimizer',
+               'train.experimental.MixedPrecisionLossScaleOptimizer'])
+class MixedPrecisionLossScaleOptimizer(optimizer.Optimizer):
+  """An optimizer that applies loss scaling.
+
+  Loss scaling is a process that multiplies the loss by a multiplier called the
+  loss scale, and divides each gradient by the same multiplier. The pseudocode
+  for this process is:
+
+  ```
+  loss = ...
+  loss *= loss_scale
+  grads = gradients(loss, vars)
+  grads /= loss_scale
+  ```
+
+  Mathematically, loss scaling has no effect, but can help avoid numerical
+  underflow in intermediate gradients when float16 tensors are used for mixed
+  precision training. By multiplying the loss, each intermediate gradient will
+  have the same multiplier applied.
+
+  The loss scale can either be a fixed constant, chosen by the user, or be
+  dynamically determined. Dynamically determining the loss scale is convenient
+  as a loss scale does not have to be explicitly chosen. However it reduces
+  performance.
+
+  This optimizer wraps another optimizer and applies loss scaling to it via a
+  `LossScale`. Loss scaling is applied whenever gradients are
+  computed, such as through `minimize()`.
+  """
+
+  def __init__(self, opt, loss_scale):
+    if not isinstance(opt, optimizer.Optimizer):
+      raise ValueError('"opt" must be an instance of Optimizer, but got: %s' %
+                       type(opt))
+    self._optimizer = opt
+
+    use_locking = opt._use_locking  # pylint: disable=protected-access
+    name = opt.get_name()
+    super(MixedPrecisionLossScaleOptimizer, self).__init__(use_locking, name)
+
+    self._loss_scale = loss_scale_module.get(loss_scale)
+    if self._loss_scale is None:
+      raise ValueError('loss_scale cannot be None')
+    self._track_trackable(self._optimizer, 'base_optimizer')
+    self._track_trackable(self._loss_scale, 'loss_scale')
+
+  def _doing_dynamic_loss_scaling(self):
+    """Check if `_loss_scale` dynamically manages the loss scale."""
+    return isinstance(self._loss_scale, loss_scale_module.DynamicLossScale)
+
+  def compute_gradients(self,
+                        loss,
+                        var_list=None,
+                        gate_gradients=optimizer.Optimizer.GATE_OP,
+                        aggregation_method=None,
+                        colocate_gradients_with_ops=False,
+                        grad_loss=None):
+    """Compute gradients of `loss` for the variables in `var_list`.
+
+    This adjusts the dynamic range of the gradient evaluation by scaling up
+    the `loss` value. The gradient values are then scaled back down by the
+    reciprocal of the loss scale. This is useful in reduced precision training
+    where small gradient values would otherwise underflow the representable
+    range.
+
+    Args:
+      loss: A Tensor containing the value to minimize or a callable taking no
+        arguments which returns the value to minimize. When eager execution is
+        enabled it must be a callable.
+      var_list: Optional list or tuple of `tf.Variable` to update to minimize
+        `loss`.  Defaults to the list of variables collected in the graph under
+        the key `GraphKeys.TRAINABLE_VARIABLES`.
+      gate_gradients: How to gate the computation of gradients.  Can be
+        `GATE_NONE`, `GATE_OP`, or `GATE_GRAPH`.
+      aggregation_method: Specifies the method used to combine gradient terms.
+        Valid values are defined in the class `AggregationMethod`.
+      colocate_gradients_with_ops: If True, try colocating gradients with the
+        corresponding op.
+      grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`.
+
+    Returns:
+      A list of (gradient, variable) pairs. Variable is always present, but
+      gradient can be `None`.
+    """
+    loss = self._scale_loss(loss)
+    grads_and_vars = self._optimizer.compute_gradients(
+        loss=loss,
+        var_list=var_list,
+        gate_gradients=gate_gradients,
+        aggregation_method=aggregation_method,
+        colocate_gradients_with_ops=colocate_gradients_with_ops,
+        grad_loss=grad_loss)
+
+    grads = [g for g, _ in grads_and_vars]
+    variables = [v for _, v in grads_and_vars]
+    unscaled_grads = self._unscale_grads(grads)
+    return list(zip(unscaled_grads, variables))
+
+  def _scale_loss(self, loss):
+    loss_scale = self._loss_scale()
+    if callable(loss):
+      def new_loss():
+        loss_val = loss()
+        return loss_val * math_ops.cast(loss_scale, loss_val.dtype)
+      return new_loss
+    else:
+      return loss * math_ops.cast(loss_scale, loss.dtype)
+
+  def _unscale_grads(self, grads):
+    loss_scale = self._loss_scale()
+    loss_scale_reciprocal = 1 / loss_scale
+    return [
+        None if g is None else self._scale_grad(g, loss_scale_reciprocal)
+        for g in grads
+    ]
+
+  def _scale_grad(self, grad, loss_scale_reciprocal):
+    if isinstance(grad, indexed_slices.IndexedSlices):
+      grad_vals = grad.values * loss_scale_reciprocal
+      return indexed_slices.IndexedSlices(grad_vals, grad.indices,
+                                          grad.dense_shape)
+    return grad * loss_scale_reciprocal
+
+  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
+    """Apply gradients to variables.
+
+    This is the second part of `minimize()`. It returns an `Operation` that
+    conditionally applies gradients if all gradient values are finite.
+    Otherwise no update is performed (nor is `global_step` incremented).
+
+    Args:
+      grads_and_vars: List of (gradient, variable) pairs as returned by
+        `compute_gradients()`.
+      global_step: Optional `Variable` to increment by one after the variables
+        have been updated.
+      name: Optional name for the returned operation.  Default to the name
+        passed to the `Optimizer` constructor.
+
+    Returns:
+      An `Operation` that conditionally applies the specified gradients. If
+      `global_step` was not None, that operation also increments `global_step`.
+
+    Raises:
+      RuntimeError: If you should use `_distributed_apply()` instead.
+    """
+    if distribute_lib.in_cross_replica_context():
+      raise ValueError('apply_gradients() must be called in a replica context.')
+
+    if not self._doing_dynamic_loss_scaling():
+      return self._optimizer.apply_gradients(grads_and_vars, global_step, name)
+
+    replica_context = distribute_lib.get_replica_context()
+    grads_and_vars = tuple(grads_and_vars)
+
+    # TODO(nluehr) cleanup GraphKeys.TRAIN_OP
+    return replica_context.merge_call(
+        self._distributed_apply, args=(grads_and_vars, global_step, name))
+
+  def _distributed_apply(self,
+                         distribution,
+                         grads_and_vars,
+                         global_step=None,
+                         name=None):
+    """A version of `apply_gradients` for cross replica context.
+
+    When users are in a cross replica strategy, they must call this rather than
+    `apply_gradients()`.
+
+    Args:
+      distribution: a `DistributionStrategy` object.
+      grads_and_vars: List of (gradient, variable) pairs as returned by
+        `compute_gradients()` and then aggregated across replicas.
+      global_step: Optional (mirrored) `Variable` to increment by one after the
+        variables have been updated.
+      name: Optional name for the returned operation. Default to the name passed
+        to the `Optimizer` constructor.
+
+    Returns:
+      An `Operation` that applies the specified gradients across all
+      replicas. If `global_step` was not None, that operation also
+      increments `global_step`
+    """
+    name = name if name is not None else self.get_name()
+    grads = [g for g, _ in grads_and_vars]
+    loss_scale_update_op, should_apply_grads = (self._loss_scale.update(grads))
+
+    def apply_fn():
+      return self._apply_gradients(distribution, grads_and_vars, global_step,
+                                   name + '-wrapped')
+
+    maybe_apply_op = smart_cond.smart_cond(should_apply_grads, apply_fn,
+                                           control_flow_ops.no_op)
+    return control_flow_ops.group(
+        maybe_apply_op, loss_scale_update_op, name=name)
+
+  def _apply_gradients(self, distribution, grads_and_vars, global_step, name):
+    """Unconditionally apply gradients in cross replica context."""
+    update_ops = distribution.extended.call_for_each_replica(
+        self._optimizer.apply_gradients,
+        args=(grads_and_vars, global_step, name))
+    return distribution.group(update_ops)
+
+  def _apply_sparse(self, grad, var):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def _apply_dense(self, grad, var):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def _resource_apply_sparse(self, grad, handle, indices):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def _resource_apply_dense(self, grad, handle):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def variables(self):
+    """Returns the variables of the Optimizer."""
+    return (self._optimizer.variables() +
+            list(self._loss_scale._weights.values()))  # pylint: disable=protected-access

tensorflow/python/training/experimental/loss_scale_optimizer_test.py

@@ -0,0 +1,315 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for MixedPrecisionLossScaleOptimizer."""
+
+import os
+
+from absl.testing import parameterized
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import custom_gradient
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import momentum
+from tensorflow.python.training.experimental import loss_scale as loss_scale_module
+from tensorflow.python.training.experimental import loss_scale_optimizer
+
+# If called outside any strategy.scope() calls, this will return the default
+# strategy.
+default_strategy_fn = distribute_lib.get_strategy
+
+
+def create_mirrored_strategy():
+  if context.num_gpus() >= 1:
+    return mirrored_strategy.MirroredStrategy(['cpu:0', 'gpu:0'])
+  else:
+    return mirrored_strategy.MirroredStrategy(['cpu:0'])
+
+
+TESTCASES = ({
+    'testcase_name': 'Base',
+    'strategy_fn': default_strategy_fn
+}, {
+    'testcase_name': 'Distribute',
+    'strategy_fn': create_mirrored_strategy
+})
+
+
+def get_gradients(opt, loss, params):
+  grads_and_vars = opt.compute_gradients(loss, params)
+  grads, _ = zip(*grads_and_vars)
+  return grads
+
+
+def create_identity_with_grad_check_fn(expected_gradient, expected_dtype=None):
+  """Returns a function that asserts it's gradient has a certain value.
+
+  This serves as a hook to assert intermediate gradients have a certain value.
+  This returns an identity function. The identity's gradient function is also
+  the identity function, except it asserts that the gradient equals
+  `expected_gradient` and has dtype `expected_dtype`.
+
+  Args:
+    expected_gradient: The gradient function asserts that the gradient is this
+      value.
+    expected_dtype: The gradient function asserts the gradient has this dtype.
+
+  Returns:
+    An identity function whose gradient function asserts the gradient has a
+    certain value.
+  """
+  @custom_gradient.custom_gradient
+  def _identity_with_grad_check(x):
+    """Function that asserts it's gradient has a certain value."""
+    x = array_ops.identity(x)
+    def grad(dx):
+      """Gradient function that asserts the gradient has a certain value."""
+      if expected_dtype:
+        assert dx.dtype == expected_dtype, (
+            'dx.dtype should be %s but is: %s' % (expected_dtype, dx.dtype))
+      expected_tensor = tensor_conversion.convert_to_tensor_v2(
+          expected_gradient, dtype=dx.dtype, name='expected_gradient'
+      )
+      # Control dependency is to ensure input is available. It's possible the
+      # dataset will throw a StopIteration to indicate there is no more data, in
+      # which case we don't want to run the assertion.
+      with ops.control_dependencies([x]):
+        assert_op = check_ops.assert_equal(dx, expected_tensor)
+      with ops.control_dependencies([assert_op]):
+        dx = array_ops.identity(dx)
+      return dx
+    return x, grad
+  # Keras sometimes has trouble serializing Lambda layers with a decorated
+  # function. So we define and return a non-decorated function.
+  def identity_with_grad_check(x):
+    return _identity_with_grad_check(x)
+  return identity_with_grad_check
+
+
+class MixedPrecisionLossScaleOptimizerTest(test.TestCase,
+                                           parameterized.TestCase):
+
+  def _run_if_in_graph_mode(self, val):
+    # Running only in graph mode is useful, because optimizers sometimes return
+    # a value that, in Graph mode, is runnable with self.evaluate. But in Eager
+    # mode, the optimizer already does the computations and the return value
+    # cannot be run.
+    if not context.executing_eagerly():
+      self.evaluate(val)
+
+  def _run_fn_with_grad_check(self, strategy, var, opt, expected_grad):
+    grad_check_fn = create_identity_with_grad_check_fn(
+        expected_grad)
+    loss = lambda: grad_check_fn(var) / strategy.num_replicas_in_sync
+    return lambda: opt.minimize(loss, var_list=[var])
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testFixedLossScaleAppliedToLossWithMinimize(self, strategy_fn):
+    with strategy_fn().scope() as strategy:
+      var = variables.Variable([5.0])
+      opt = gradient_descent.GradientDescentOptimizer(2.0)
+      loss_scale = 10.
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      # We need num_replicas_in_sync to divide loss_scale, otherwise loss_scale
+      # / strategy.num_replicas_in_sync will not be exact, which could lead to
+      # assertion failures due to rounding issues.
+      self.assertEqual(loss_scale % strategy.num_replicas_in_sync, 0)
+      run_fn = self._run_fn_with_grad_check(
+          strategy, var, opt, loss_scale / strategy.num_replicas_in_sync)
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # The loss is the identity of the variable. Therefore the gradient is 1,
+      # and so the variable will be init_val - grad * lr == 5 - 1 * 2 == 3
+      self.assertAllClose([3.], self.evaluate(var))
+
+  @test_util.deprecated_graph_mode_only
+  def testFixedLossScaleAppliedToLossWithGetGradients(self):
+    var = variables.Variable([2.0])
+    opt = gradient_descent.GradientDescentOptimizer(1.0)
+    loss_scale = 10.
+    opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(opt, loss_scale)
+    grad_check_fn = create_identity_with_grad_check_fn(loss_scale)
+    loss = grad_check_fn(var)
+    run_op = get_gradients(opt, loss, [var])
+    self.evaluate(variables.global_variables_initializer())
+    # This will cause an assertion to run, as
+    # create_identity_with_grad_check_fn added an assertion op.
+    self.evaluate(run_op)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testDynamicLossScale(self, strategy_fn):
+    strategy = strategy_fn()
+    learning_rate = 2.
+    expected_gradient = resource_variable_ops.ResourceVariable(
+        learning_rate / strategy.num_replicas_in_sync)
+    with strategy.scope():
+      var = variables.Variable([5.0])
+      opt = gradient_descent.GradientDescentOptimizer(learning_rate)
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=2, increment_period=1, multiplier=2)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      self.assertEqual(
+          loss_scale.initial_loss_scale % strategy.num_replicas_in_sync, 0)
+
+      run_fn = self._run_fn_with_grad_check(strategy, var, opt,
+                                            expected_gradient)
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # The loss is the identity of the variable. Therefore the gradient is 1,
+      # and so the variable will be init_val - grad * lr == 5 - 1 * 2 == 3
+      self.assertAllClose([3.], self.evaluate(var))
+
+      # Loss scale will be double, so the expected gradient is also doubled.
+      self.evaluate(
+          expected_gradient.assign(2 * learning_rate /
+                                   strategy.num_replicas_in_sync))
+      run_op = strategy.experimental_run(run_fn)
+      self._run_if_in_graph_mode(run_op)
+      # As before, the 2 is subtracted from the variable, making it's new value
+      # 1.
+      self.assertAllClose([1.], self.evaluate(var))
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testDynamicUpdate(self, strategy_fn):
+    with strategy_fn().scope() as strategy:
+      var = variables.Variable([1.0, 2.0])
+      opt = gradient_descent.GradientDescentOptimizer(1.0)
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=2, increment_period=1, multiplier=2)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+
+      # Test optimizer with finite gradients
+      loss = lambda: var * 2.0 / strategy.num_replicas_in_sync
+      run_fn = lambda: opt.minimize(loss, var_list=[var])
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # Gradient is 2, so variable will have 2 subtracted from it
+      self.assertAllClose([-1.0, 0.0], self.evaluate(var))
+      # Loss scale has doubled from 2 to 4
+      self.assertEqual(4., self.evaluate(opt._loss_scale()))
+
+      # Test optimizer with NaN gradients
+      loss = lambda: var * float('NaN')
+      run_fn = lambda: opt.minimize(loss, var_list=[var])
+      run_op = strategy.experimental_run(run_fn)
+      self._run_if_in_graph_mode(run_op)
+      # Variable should not change from before, due to NaN gradients.
+      self.assertAllClose(self.evaluate(var), [-1.0, 0.0])
+      # Loss scale should half due to NaN gradients.
+      self.assertEqual(2., self.evaluate(opt._loss_scale()))
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testDynamicLossScaleWithSlots(self, strategy_fn):
+    with strategy_fn().scope() as strategy:
+      var = variables.Variable([1.0, 2.0])
+      # An SGD optimizer with momentum has slot variables.
+      opt = momentum.MomentumOptimizer(1.0, momentum=1.)
+      initial_loss_scale = 2.
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=initial_loss_scale,
+          increment_period=1,
+          multiplier=4)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      loss = lambda: var / strategy.num_replicas_in_sync
+      run_fn = lambda: opt.minimize(loss, var_list=[var])
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # The momentum accumulator starts at 0 and the gradient is 1. The
+      # accumulator is incremented by the gradient, so it is now 1. Then the
+      # variable is subtracted by the accumulator, so the variable is subtracted
+      # by 1.
+      self.assertAllClose([0.0, 1.0], self.evaluate(var))
+      self.assertEqual(self.evaluate(opt._loss_scale()), initial_loss_scale * 4)
+
+      run_op = strategy.experimental_run(run_fn)
+      self._run_if_in_graph_mode(run_op)
+      # The momentum accumulator was 1 before this step and the gradient is 1.
+      # The accumulator is incremented by the gradient, so it is now 2. Then the
+      # variable is subtracted by the accumulator, so the variable is subtracted
+      # by 2.
+      self.assertAllClose([-2., -1.], self.evaluate(var))
+      self.assertEqual(
+          self.evaluate(opt._loss_scale()), initial_loss_scale * 16)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testCheckpoint(self, strategy_fn):
+    strategy = strategy_fn()
+    if (isinstance(strategy, mirrored_strategy.MirroredStrategy) and
+        not context.executing_eagerly()):
+      # TODO(b/121381184): Enable running the test in this case.
+      return
+
+    with self.test_session(), strategy.scope():
+      # Build and run a simple model.
+      var = variables.Variable([2.0])
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=1., increment_period=2., multiplier=2.)
+      opt = momentum.MomentumOptimizer(1.0, momentum=1.)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      run_fn = lambda: opt.minimize(lambda: var + 1., var_list=[var])
+      opt_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self.evaluate(opt_op)
+      self.assertEqual(self.evaluate(loss_scale()), 1.)
+      self.assertEqual(self.evaluate(loss_scale._num_good_steps), 1)
+
+      # Save a checkpoint.
+      checkpoint = trackable_utils.Checkpoint(optimizer=opt)
+      prefix = os.path.join(self.get_temp_dir(), 'ckpt')
+      save_path = checkpoint.save(prefix)
+
+      # Run model again.
+      self.evaluate(strategy.experimental_run(run_fn))
+      self.assertEqual(self.evaluate(loss_scale()), 2.)
+      self.assertEqual(self.evaluate(loss_scale._num_good_steps), 0)
+
+      # Load checkpoint and ensure loss scale is back to it's original value.
+      status = checkpoint.restore(save_path)
+      status.assert_consumed()
+      status.run_restore_ops()
+      self.assertEqual(self.evaluate(loss_scale()), 1.)
+      self.assertEqual(self.evaluate(loss_scale._num_good_steps), 1)
+
+  def testPassingNoneToLossScale(self):
+    opt = gradient_descent.GradientDescentOptimizer(1.0)
+    with self.assertRaisesRegex(ValueError, r'loss_scale cannot be None'):
+      loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(opt, None)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/experimental/loss_scale_test.py

@@ -0,0 +1,325 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for LossScale classes.."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training.experimental import loss_scale as loss_scale_module
+
+# TODO(reedwm): Create test case using multiple graphs
+
+# If called outside any strategy.scope() calls, this will return the default
+# strategy.
+default_strategy_fn = distribute_lib.get_strategy
+
+
+def create_mirrored_strategy():
+  if context.num_gpus() >= 1:
+    return mirrored_strategy.MirroredStrategy(['cpu:0', 'gpu:0'])
+  else:
+    return mirrored_strategy.MirroredStrategy(['cpu:0'])
+
+
+TESTCASES = ({
+    'testcase_name': 'base',
+    'strategy_fn': default_strategy_fn
+}, {
+    'testcase_name': 'distribute',
+    'strategy_fn': create_mirrored_strategy
+})
+
+
+class FixedLossScaleTest(test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_basic(self):
+    loss_scale_value = 1000
+    loss_scale = loss_scale_module.FixedLossScale(loss_scale_value)
+
+    update_op, should_apply = loss_scale.update([constant_op.constant(0.)])
+    self.evaluate(update_op)
+    # should_apply should be a bool instead of a tensor, so that a tf.cond does
+    # not have to be built in the graph by the caller.
+    self.assertIsInstance(should_apply, bool)
+    self.assertTrue(should_apply)
+    self.assertEqual(loss_scale_value, self.evaluate(loss_scale()))
+
+    update_op, should_apply = loss_scale.update(
+        [constant_op.constant(float('NaN'))])
+    self.evaluate(update_op)
+    self.assertIsInstance(should_apply, bool)
+    self.assertTrue(should_apply)
+    self.assertEqual(loss_scale_value, self.evaluate(loss_scale()))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_serialization(self):
+    loss_scale = loss_scale_module.get(123)
+    config = loss_scale.get_config()
+    loss_scale = loss_scale_module.FixedLossScale.from_config(config)
+    self.assertEqual(self.evaluate(loss_scale()), 123.)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_call_type(self):
+    scalar = loss_scale_module.FixedLossScale(123)
+    self.assertIsInstance(scalar(), tensor_lib.Tensor)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_repr(self):
+    loss_scale = loss_scale_module.FixedLossScale(123)
+    self.assertEqual(repr(loss_scale), 'FixedLossScale(123.0)')
+
+
+def _get_example_iter(inputs):
+  dataset = dataset_ops.Dataset.from_tensor_slices(inputs)
+  return dataset_ops.make_one_shot_iterator(dataset)
+
+
+class DynamicLossScaleTest(test.TestCase, parameterized.TestCase):
+
+  def _get_tensor(self, is_finite):
+    tensor = cond.cond(is_finite, lambda: 1., lambda: float('NaN'))
+
+    if not distribute_lib.has_strategy():
+      return tensor
+
+    def get():
+      rep_id = (
+          distribute_lib.get_replica_context()
+          .replica_id_in_sync_group)
+      return cond.cond(
+          math_ops.equal(rep_id, 0), lambda: tensor, lambda: 1.)
+
+    distribution = distribute_lib.get_strategy()
+    return distribution.extended.call_for_each_replica(get)
+
+  def _test_helper(self,
+                   inputs,
+                   expected_outputs,
+                   initial_loss_scale=1.,
+                   increment_period=2,
+                   multiplier=2):
+    loss_scale = loss_scale_module.DynamicLossScale(
+        initial_loss_scale=initial_loss_scale,
+        increment_period=increment_period,
+        multiplier=multiplier)
+    itr = _get_example_iter(inputs)
+
+    def update():
+      is_finite = itr.get_next()
+      grad = self._get_tensor(is_finite)
+      update_op, should_apply_gradients = loss_scale.update([grad])
+      assert_op = check_ops.assert_equal(should_apply_gradients, is_finite)
+      if context.executing_eagerly():
+        return
+      with ops.control_dependencies([assert_op]):
+        return array_ops.identity(update_op)
+
+    actual_outputs = []
+
+    if not context.executing_eagerly():
+      update_op = update()
+      self.evaluate(variables.global_variables_initializer())
+    for _ in range(len(inputs)):
+      if context.executing_eagerly():
+        update()
+      else:
+        self.evaluate(update_op)
+      actual_outputs.append(self.evaluate(loss_scale()))
+    self.assertEqual(actual_outputs, expected_outputs)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_increase(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [True] * 6
+      expected_outputs = [1, 2, 2, 4, 4, 8]
+      self._test_helper(inputs, expected_outputs)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_keep_increasing_until_capped(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = np.finfo(np.float32).max / 4
+      max_float = np.finfo(np.float32).max
+
+      inputs = [True] * 6
+      # Output is capped the 2nd time it doubles.
+      expected_outputs = [
+          init_loss_scale, init_loss_scale * 2, init_loss_scale * 2, max_float,
+          max_float, max_float
+      ]
+
+      self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_decrease_every_step(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [False] * 6
+      init_loss_scale = 1024
+      expected_outputs = [512, 256, 128, 64, 32, 16]
+
+    self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_keep_decreasing_until_one(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [False] * 6
+      init_loss_scale = 16
+      expected_outputs = [8, 4, 2, 1, 1, 1]
+
+      self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_nan_clear_good_step(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [True, True, True, False, True]
+      expected_outputs = [1, 2, 2, 1, 1]
+      self._test_helper(inputs, expected_outputs)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_trigger_loss_scale_update_each_step(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 1
+      increment_period = 1
+
+      inputs = [True] * 3 + [False, True, True]
+      expected_outputs = [2, 4, 8, 4, 8, 16]
+
+      self._test_helper(inputs, expected_outputs, init_loss_scale,
+                        increment_period)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_alternating_good_and_bad_gradients_trigger_each_step(
+      self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 1
+      increment_period = 1
+
+      inputs = [True, False] * 4 + [True]
+      expected_outputs = [2, 1, 2, 1, 2, 1, 2, 1, 2]
+      self._test_helper(inputs, expected_outputs, init_loss_scale,
+                        increment_period)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_alternating_good_and_bad_gradients_trigger_every_other_step(
+      self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 32
+      increment_period = 2
+
+      inputs = [True, False] * 3 + [True]
+      expected_outputs = [32, 16, 16, 8, 8, 4, 4]
+      self._test_helper(inputs, expected_outputs, init_loss_scale,
+                        increment_period)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_nondefault_multiplier(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 4
+      multiplier = 3
+      inputs = [True, True, False, True, True]
+      expected_outputs = [4, 12, 4, 4, 12]
+      self._test_helper(
+          inputs, expected_outputs, init_loss_scale, multiplier=multiplier)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_random_mix_good_and_bad_gradients(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 4
+      inputs = [
+          False, True, True, True, False, True, False, True, True, True, False
+      ]
+      expected_outputs = [2, 2, 4, 4, 2, 2, 1, 1, 2, 2, 1]
+      self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_single_tensor_gradient(self, strategy_fn):
+    with strategy_fn().scope():
+      loss_scale = loss_scale_module.DynamicLossScale()
+      grad = constant_op.constant(4.0)
+      _, should_apply = loss_scale.update(grad)
+      self.assertTrue(self.evaluate(should_apply))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_serialization(self):
+    loss_scale = loss_scale_module.DynamicLossScale(
+        initial_loss_scale=1, increment_period=2, multiplier=3)
+    config = loss_scale.get_config()
+    loss_scale = loss_scale_module.DynamicLossScale.from_config(config)
+    self.evaluate(variables.global_variables_initializer())
+    self.assertEqual(self.evaluate(loss_scale()), 1)
+    self.assertEqual(loss_scale.increment_period, 2)
+    self.assertEqual(loss_scale.multiplier, 3)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_update_with_none_gradients(self):
+    loss_scale = loss_scale_module.DynamicLossScale()
+    loss_scale.update([None])
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_get(self):
+    scalar = loss_scale_module.get('dynamic')
+    scalar2 = loss_scale_module.DynamicLossScale()
+    self.assertEqual(scalar.initial_loss_scale, scalar2.initial_loss_scale)
+    self.assertEqual(scalar.increment_period, scalar2.increment_period)
+    self.assertEqual(scalar.multiplier, scalar2.multiplier)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_call_type(self):
+    scalar = loss_scale_module.DynamicLossScale()
+    self.assertIsInstance(scalar(), tensor_lib.Tensor)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_repr(self, strategy_fn):
+    with strategy_fn().scope():
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=1, increment_period=2, multiplier=3)
+      if context.executing_eagerly():
+        self.assertEqual(repr(loss_scale),
+                         'DynamicLossScale(current_loss_scale=1.0, '
+                         'num_good_steps=0, initial_loss_scale=1.0, '
+                         'increment_period=2, multiplier=3.0)')
+      else:
+        self.assertEqual(repr(loss_scale),
+                         'DynamicLossScale(initial_loss_scale=1.0, '
+                         'increment_period=2, multiplier=3.0)')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/experimental/mixed_precision.py

@@ -0,0 +1,248 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains functions to use mixed precision with the graph rewrite."""
+
+from tensorflow.python.framework import config
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.training import optimizer
+from tensorflow.python.training.experimental import loss_scale_optimizer as loss_scale_optimizer_v1
+from tensorflow.python.training.experimental import mixed_precision_global_state
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+# A mapping between optimizers and (wrapper_fn, wrapper_cls) pairs. wrapper_cls
+# is a loss scale optimizer class, and wrapper_fn is a function that takes in
+# an optimizer and LossScale and returns a wrapper_cls instance.
+_REGISTERED_WRAPPER_OPTIMIZER_CLS = {
+    optimizer.Optimizer:
+        (loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer,) * 2,
+}
+
+
+@tf_export('__internal__.mixed_precision.register_loss_scale_wrapper', v1=[])
+def register_loss_scale_wrapper(optimizer_cls, wrapper_fn, wrapper_cls=None):
+  """Registers a loss scale optimizer wrapper.
+
+  `tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite`
+  automatically wraps an optimizer with an optimizer wrapper that performs loss
+  scaling. This function registers a
+  `(base_cls, wrapper_fn, wrapper_cls)` triple
+  that is used by `enable_mixed_precision_graph_rewrite`, where
+  `wrapper_fn` is called to create a `wrapper_cls` instance that wraps an
+  `optimizer_cls` instance.
+
+  Args:
+    optimizer_cls: A base optimizer class, e.g. `tf.keras.optimizers.Optimizer`.
+    wrapper_fn: A function that takes in arguments "optimizer" and
+      "loss_scale", and returns a loss scale optimizer of type "wrapper_cls"
+      that wraps "optimizer".
+    wrapper_cls: A loss scale optimizer class. Defaults to `wrapper_fn`, in
+      which case `wrapper_fn` should be a loss scale optimizer class whose
+      constructor takes in arguments "optimizer" and "loss_scale".
+  """
+  _REGISTERED_WRAPPER_OPTIMIZER_CLS[optimizer_cls] = (
+      wrapper_fn, wrapper_cls or wrapper_fn)
+
+
+def _wrap_optimizer(opt, loss_scale):
+  """Wraps an optimizer with a LossScaleOptimizer."""
+
+  for _, wrapper_optimizer in _REGISTERED_WRAPPER_OPTIMIZER_CLS.values():
+    if isinstance(opt, wrapper_optimizer):
+      raise ValueError('"opt" must not already be an instance of a {cls}. '
+                       '`enable_mixed_precision_graph_rewrite` will '
+                       'automatically wrap the optimizer with a '
+                       '{cls}.'
+                       .format(cls=wrapper_optimizer.__name__))
+
+  for optimizer_cls, (wrapper_fn, _) in (
+      _REGISTERED_WRAPPER_OPTIMIZER_CLS.items()):
+    if isinstance(opt, optimizer_cls):
+      return wrapper_fn(opt, loss_scale)
+
+  raise ValueError('"opt" must be an instance of a tf.train.Optimizer or a '
+                   'tf.keras.optimizers.Optimizer, but got: %s' % opt)
+
+
+@deprecation.deprecated_endpoints(
+    'train.experimental.enable_mixed_precision_graph_rewrite')
+@tf_export(v1=['mixed_precision.enable_mixed_precision_graph_rewrite',
+               'train.experimental.enable_mixed_precision_graph_rewrite'])
+def enable_mixed_precision_graph_rewrite_v1(opt, loss_scale='dynamic'):
+  """Enable mixed precision via a graph rewrite.
+
+  Mixed precision is the use of both float32 and float16 data types when
+  training a model to improve performance. This is achieved via a graph rewrite
+  operation and a loss-scale optimizer.
+
+  Performing arithmetic operations in float16 takes advantage of specialized
+  processing units, such as NVIDIA Tensor Cores, for much higher arithmetic
+  throughput. However, due to the smaller representable range, performing the
+  entire training with float16 can result in gradient underflow, that is, small
+  gradient values becoming zeroes. Instead, performing only select arithmetic
+  operations in float16 results in higher throughput and decreased training
+  time when using compatible hardware accelerators while also reducing memory
+  usage, typically without sacrificing model accuracy.
+
+  Note: While the mixed precision rewrite changes the datatype of various
+  layers throughout the model, the same accuracy reached in float32 is
+  expected. If a `NaN` gradient occurs with dynamic loss scaling, the model
+  update for that batch is skipped. In this case, the global step count is not
+  incremented, and the `LossScaleOptimizer` attempts to decrease the loss
+  scaling value to avoid `NaN` values in subsequent iterations. This approach
+  has been shown to achieve the same accuracy as float32 and, in most cases,
+  better training throughput.
+
+  Example:
+
+  ```python
+  model = tf.keras.models.Sequential([
+      tf.keras.layers.Dense(64, activation='relu'),
+      tf.keras.layers.Dense(64, activation='softmax'),
+  ])
+
+  opt = tf.keras.optimizers.SGD()
+  opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
+  model.compile(loss="mse", optimizer=opt)
+
+  x_train = np.random.random((1024, 64))
+  y_train = np.random.random((1024, 64))
+  model.fit(x_train, y_train)
+  ```
+
+  Calling `enable_mixed_precision_graph_rewrite(opt)` enables the graph rewrite
+  operation before computing gradients. The function additionally returns an
+  `Optimizer` (`opt`) wrapped with a `LossScaleOptimizer`. This prevents
+  underflow in the float16 tensors during the backward pass. An optimizer of
+  type `tf.train.Optimizer` or `tf.keras.optimizers.Optimizer` must be passed
+  to this function, which will then be wrapped to use loss scaling.
+
+  The graph rewrite operation changes the `dtype` of certain operations in the
+  graph from float32 to float16. There are several categories of operations
+  that are either included or excluded by this rewrite operation. The following
+  categories of Ops are defined inside corresponding functions under the class
+  `AutoMixedPrecisionLists` in
+  <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/
+  core/grappler/optimizers/auto_mixed_precision_lists.h">
+  auto_mixed_precision_lists.h</a>:
+
+  * `ClearList`: Ops that do not have numerically significant adverse effects.
+  E.g. `ArgMax` and `Floor`.
+  * `AllowList`: Ops that are considered numerically safe for execution in
+  float16, and thus are always converted. E.g. `Conv2D`.
+  * `DenyList`: Ops that are numerically unsafe to execute in float16 and
+  can negatively affect downstream nodes. E.g. `Softmax`.
+  * `GrayList`: Ops that are considered numerically safe for execution in
+  float16 unless downstream from a DenyList Op. E.g. `Add` and `AvgPool`.
+
+  When this function is used, gradients should only be computed and applied
+  with the returned optimizer, either by calling `opt.minimize()` or
+  `opt.compute_gradients()` followed by `opt.apply_gradients()`.
+  Gradients should not be computed with `tf.gradients` or `tf.GradientTape`.
+  This is because the returned optimizer will apply loss scaling, and
+  `tf.gradients` or `tf.GradientTape` will not. If you do directly use
+  `tf.gradients` or `tf.GradientTape`, your model may not converge due to
+  float16 underflow problems.
+
+  When eager execution is enabled, the mixed precision graph rewrite is only
+  enabled within `tf.function`s, as outside `tf.function`s, there is no graph.
+
+  For NVIDIA GPUs with Tensor cores, as a general performance guide, dimensions
+  (such as batch size, input size, output size, and channel counts)
+  should be powers of two if under 256, or  otherwise divisible by 8 if above
+  256. For more information, check out the
+  [NVIDIA Deep Learning Performance Guide](
+  https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html).
+
+  Currently, mixed precision is only enabled on NVIDIA Tensor Core GPUs with
+  Compute Capability 7.0 and above (Volta, Turing, or newer architectures). The
+  parts of the graph on CPUs and TPUs are untouched by the graph rewrite.
+
+  Raises:
+    `ValueError`, if the `tf.keras.mixed_precision` API is also used by calling
+    `tf.keras.mixed_precision.set_global_policy`. Only one mixed precision
+    API can be used.
+
+  Args:
+    opt: An instance of a `tf.keras.optimizers.Optimizer` or a
+      `tf.train.Optimizer`.
+    loss_scale: Either an int/float, the string `"dynamic"`, or an instance of
+      a `tf.mixed_precision.experimental.LossScale`. The loss scale to use. It
+      is recommended to keep this as its default value of `"dynamic"`, which
+      will adjust the scaling automatically to prevent `Inf` or `NaN` values.
+
+  Returns:
+    A version of `opt` that will use loss scaling to prevent underflow.
+  """
+  if mixed_precision_global_state.is_using_mixed_precision_policy():
+    raise ValueError(
+        'The mixed precision graph rewrite cannot be enabled, because the '
+        'global Keras dtype Policy has been set to a mixed precision policy. '
+        'At most, one of the following can be called:\n\n'
+        '  1. tf.keras.mixed_precision.set_global_policy() with a mixed '
+        'precision policy (You called this first)\n\n'
+        '  2. tf.train.experimental.enable_mixed_precision_graph_rewrite() '
+        '(You called this second)\n'
+        'You called both functions, which is an error, because both functions '
+        'enable you to use mixed precision. If in doubt which function to use, '
+        'use the first, as it supports Eager execution and is more '
+        'customizable.')
+
+  if mixed_precision_global_state.non_mixed_precision_session_created():
+    # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the
+    # Sessions have already been closed, do not raise this error message.
+    tf_logging.warn('You already have existing Sessions that do not use mixed '
+                    'precision. enable_mixed_precision_graph_rewrite() will '
+                    'not affect these Sessions.')
+  opt = _wrap_optimizer(opt, loss_scale)
+  config.set_optimizer_experimental_options({'auto_mixed_precision': True})
+  mixed_precision_global_state.set_mixed_precision_graph_rewrite_enabled(True)
+  return opt
+
+
+@deprecation.deprecated_endpoints(
+    'train.experimental.disable_mixed_precision_graph_rewrite')
+@tf_export(v1=['mixed_precision.disable_mixed_precision_graph_rewrite',
+               'train.experimental.disable_mixed_precision_graph_rewrite'])
+def disable_mixed_precision_graph_rewrite_v1():
+  """Disables the mixed precision graph rewrite.
+
+  After this is called, the mixed precision graph rewrite will no longer run for
+  new Sessions, and so float32 operations will no longer be converted to float16
+  in such Sessions. However, any existing Sessions will continue to have the
+  graph rewrite enabled if they were created after
+  `enable_mixed_precision_graph_rewrite` was called but before
+  `disable_mixed_precision_graph_rewrite` was called.
+
+  This does not undo the effects of loss scaling. Any optimizers wrapped with a
+  LossScaleOptimizer will continue to do loss scaling, although this loss
+  scaling will no longer be useful if the optimizer is used in new Sessions, as
+  the graph rewrite no longer converts the graph to use float16.
+
+  This function is useful for unit testing. A unit tests can test using the
+  mixed precision graph rewrite, then disable it so future unit tests continue
+  using float32. If this is done, unit tests should not share a single session,
+  as `enable_mixed_precision_graph_rewrite` and
+  `disable_mixed_precision_graph_rewrite` have no effect on existing sessions.
+  """
+  # We only have a separate V1 version of this function, because the V1
+  # docstring mentions sessions.
+  if (not
+      mixed_precision_global_state.is_mixed_precision_graph_rewrite_enabled()):
+    tf_logging.warn('disable_mixed_precision_graph_rewrite() called when mixed '
+                    'precision is already disabled.')
+  config.set_optimizer_experimental_options({'auto_mixed_precision': False})
+  mixed_precision_global_state.set_mixed_precision_graph_rewrite_enabled(False)

tensorflow/python/training/experimental/mixed_precision_global_state.py

@@ -0,0 +1,66 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains global variables related to mixed precision.
+
+This is not part of mixed_precision.py to avoid a circular dependency.
+mixed_precision.py depends on Session, and Session depends on this file.
+"""
+
+from tensorflow.python.util.tf_export import tf_export
+
+# Whether the mixed precision graph rewrite has been enabled or not with
+# `enable_mixed_precision_graph_rewrite`. Used to turn on auto_mixed_precision
+# in ConfigProtos passed to Sessions.
+_mixed_precision_graph_rewrite_is_enabled = False
+
+
+# True if a Session has been created without the mixed precision graph rewrite
+# being enabled. Used to give a warning if mixed precision is enabled after a
+# Session has already been created.
+_non_mixed_precision_session_created = False
+
+# Whether the global tf.keras.mixed_precision.Policy uses mixed precision. Used
+# to raise an error message if both a mixed Policy and the graph rewrite are
+# used at the same time.
+_using_mixed_precision_policy = False
+
+
+@tf_export('__internal__.train.is_mixed_precision_graph_rewrite_enabled', v1=[])
+def is_mixed_precision_graph_rewrite_enabled():
+  return _mixed_precision_graph_rewrite_is_enabled
+
+
+def set_mixed_precision_graph_rewrite_enabled(enabled):
+  global _mixed_precision_graph_rewrite_is_enabled
+  _mixed_precision_graph_rewrite_is_enabled = enabled
+
+
+def non_mixed_precision_session_created():
+  return _non_mixed_precision_session_created
+
+
+def set_non_mixed_precision_session_created(created):
+  global _non_mixed_precision_session_created
+  _non_mixed_precision_session_created = created
+
+
+def is_using_mixed_precision_policy():
+  return _using_mixed_precision_policy
+
+
+@tf_export('__internal__.train.set_using_mixed_precision_policy', v1=[])
+def set_using_mixed_precision_policy(is_using):
+  global _using_mixed_precision_policy
+  _using_mixed_precision_policy = is_using

tensorflow/python/training/experimental/mixed_precision_test.py

@@ -0,0 +1,207 @@
+# Copyright 2019 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.
+# ==============================================================================
+
+import os
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.training import gradient_descent as gradient_descent_v1
+from tensorflow.python.training.experimental import loss_scale_optimizer as loss_scale_optimizer_v1
+from tensorflow.python.training.experimental import mixed_precision
+from tensorflow.python.training.experimental import mixed_precision_global_state
+
+
+class MixedPrecisionTest(test.TestCase, parameterized.TestCase):
+
+  IGNORE_PERF_VAR = 'TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'
+
+  def setUp(self):
+    super(MixedPrecisionTest, self).setUp()
+    # Enable the tests to be run on pre-Volta GPUs by telling the grappler pass
+    # to ignore performance and always transform the graph.
+    self._original_ignore_perf_value = os.getenv(self.IGNORE_PERF_VAR)
+    os.environ[self.IGNORE_PERF_VAR] = '1'
+
+  def tearDown(self):
+    # Set the IGNORE_PERF_VAR variable back to it's original value.
+    if self._original_ignore_perf_value is not None:
+      os.environ[self.IGNORE_PERF_VAR] = self._original_ignore_perf_value
+    else:
+      del os.environ[self.IGNORE_PERF_VAR]
+
+    mixed_precision.disable_mixed_precision_graph_rewrite_v1()
+    super(MixedPrecisionTest, self).tearDown()
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_wrap_optimizer(self):
+    opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
+    opt = mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+    self.assertIsInstance(
+        opt, loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer)
+    self.assertEqual(self.evaluate(opt._loss_scale()), 123.)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_optimizer_errors(self):
+    opt = 1
+    expected_regex = ('"opt" must be an instance of a tf.train.Optimizer or '
+                      'a tf.keras.optimizers.Optimizer, but got')
+    with self.assertRaisesRegex(ValueError, expected_regex):
+      mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt)
+    self.assertFalse(config.get_optimizer_experimental_options()
+                     .get('auto_mixed_precision', False))
+
+    opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
+    opt = loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer(opt,
+                                                                   'dynamic')
+    with self.assertRaisesRegex(
+        ValueError, '"opt" must not already be an instance of a '
+        'MixedPrecisionLossScaleOptimizer.'):
+      mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt)
+    self.assertFalse(config.get_optimizer_experimental_options()
+                     .get('auto_mixed_precision', False))
+
+  @test_util.run_in_graph_and_eager_modes()
+  def test_register_loss_scale_wrapper_with_2_arguments(self):
+    class MyOptimizer:
+      pass
+
+    class MyLossScaleOptimizer(MyOptimizer):
+
+      def __init__(self, inner_optimizer, loss_scale):
+        self.inner_optimizer = inner_optimizer
+        self.loss_scale = loss_scale
+
+    mixed_precision.register_loss_scale_wrapper(MyOptimizer,
+                                                MyLossScaleOptimizer)
+    opt = MyOptimizer()
+    opt = mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+    self.assertIsInstance(opt, MyLossScaleOptimizer)
+    self.assertEqual(opt.loss_scale, 123.)
+
+  @test_util.run_in_graph_and_eager_modes()
+  def test_register_loss_scale_wrapper_with_3_arguments(self):
+    class MyOptimizer:
+      pass
+
+    class MyLossScaleOptimizer(MyOptimizer):
+
+      def __init__(self, inner_optimizer, loss_scale):
+        self.inner_optimizer = inner_optimizer
+        self.loss_scale = loss_scale
+
+    is_called = False
+
+    def create_lso(inner_optimizer, loss_scale):
+      nonlocal is_called
+      is_called = True
+      return MyLossScaleOptimizer(inner_optimizer, loss_scale)
+
+    mixed_precision.register_loss_scale_wrapper(MyOptimizer,
+                                                create_lso,
+                                                MyLossScaleOptimizer)
+    opt = MyOptimizer()
+    opt = mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+    self.assertIsInstance(opt, MyLossScaleOptimizer)
+    self.assertEqual(opt.loss_scale, 123.)
+    self.assertTrue(is_called)
+
+  @test_util.run_gpu_only
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_tfrt('Grappler rewrite doesn\'t apply to tfrt.')
+  def test_grappler_pass_enabled(self):
+    opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
+    mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+
+    var = variables.Variable([[1.0]])
+
+    def overflow_in_float16():
+      out = var * 2 ** 10
+      out = math_ops.matmul(out, out)
+      return array_ops.reshape(out, ())
+
+    if context.executing_eagerly():
+      f = def_function.function(overflow_in_float16)
+      self.assertEqual(f().numpy(), float('Inf'))
+      # Outside a def_function.function, the grappler pass will not be applied.
+      self.assertAlmostEqual(overflow_in_float16().numpy(), 2 ** 20)
+
+      # Test disabling mixed precision.
+      mixed_precision.disable_mixed_precision_graph_rewrite_v1()
+      self.assertEqual(f().numpy(), 2 ** 20)
+    else:
+      with session.Session() as sess:
+        out = overflow_in_float16()
+        sess.run(var.initializer)
+        self.assertEqual(sess.run(out), float('Inf'))
+
+      # Test Session will enable the auto_mixed_precision grappler pass in a
+      # ConfigProto passed by the user
+      with session.Session(config=config_pb2.ConfigProto()) as sess:
+        out = overflow_in_float16()
+        sess.run(var.initializer)
+        self.assertEqual(sess.run(out), float('Inf'))
+
+      # Test disabling mixed precision.
+      mixed_precision.disable_mixed_precision_graph_rewrite_v1()
+      with session.Session() as sess:
+        out = overflow_in_float16()
+        sess.run(var.initializer)
+        self.assertAlmostEqual(sess.run(out), 2 ** 20)
+
+  @test.mock.patch.object(tf_logging, 'warn')
+  def test_warn_if_session_already_exists(self, mock_warn):
+    # Set this to False, so Sessions created in previous tests do not trigger
+    # the warning.
+    mixed_precision_global_state.set_non_mixed_precision_session_created(False)
+
+    with session.Session():
+      mixed_precision.enable_mixed_precision_graph_rewrite_v1(
+          gradient_descent_v1.GradientDescentOptimizer(1.0))
+      mock_warn.assert_any_call(
+          'You already have existing Sessions that do not use mixed precision. '
+          'enable_mixed_precision_graph_rewrite() will not affect these '
+          'Sessions.')
+
+  @test.mock.patch.object(tf_logging, 'warn')
+  def test_do_not_warn_if_session_does_not_already_exist(self, mock_warn):
+    # Set this to False, so Sessions created in previous tests do not trigger
+    # the warning.
+    mixed_precision_global_state.set_non_mixed_precision_session_created(False)
+
+    mixed_precision.enable_mixed_precision_graph_rewrite_v1(
+        gradient_descent_v1.GradientDescentOptimizer(1.0))
+    with session.Session():
+      # Make sure the "You already have existing Sessions" warning was not
+      # issued, since the Session was only created after
+      # enable_mixed_precision_graph_rewrite.
+      for call_arg in mock_warn.call_args_list:
+        msg = call_arg[0][0]
+        self.assertNotIn('You already have existing Sessions that do not use '
+                         'mixed precision', msg)
+
+
+if __name__ == '__main__':
+  test.main()