code: 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
  
file_path: tensorflow-upstream/tensorflow/python/training/experimental/mixed_precision.py
imports: 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
  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
label_desc: Use image node_id 3 for calling a global function with example usage: enable_mixed_precision_graph_rewrite_v1(opt, loss_scale) and returns: opt