Merge the two LossScale classes.

Before, we had the following loss scale base classes, which did the exact same thing:
  * tf.keras.mixed_precision.experimental.LossScale, which only worked for the keras OptimizerV2
  * An unexposed LossScale in tensorflow/python/training/experimental, which only worked for the V1 Optimizer

This change removes the Keras LossScale and merges it into the training LossScale, which now works for both Optimizers. The training LossScale is exposed as tf.train.experimental.LossScale. I moved over some functionality, comments, and style conventions from the Keras LossScale to the training LossScale.

Because the LossScale class can not rely on Keras, the Keras OptimizerV2 now calls backend.track_variable on the LossScale variables instead.

Note: I intend to cherrypick this into TF 1.14.
PiperOrigin-RevId: 248213961

Author: reedwm

tensorflow/python/keras/mixed_precision/experimental/BUILD

@@ -96,7 +96,7 @@ py_library(
     srcs = ["loss_scale_optimizer.py"],
     srcs_version = "PY2AND3",
     deps = [
-        ":loss_scale",
+        "//tensorflow/python:loss_scale",
         "//tensorflow/python/keras/optimizer_v2",
         "@absl_py//absl/testing:parameterized",
     ],
@@ -116,31 +116,6 @@ cuda_py_test(
     ],
 )
 
-py_library(
-    name = "loss_scale",
-    srcs = ["loss_scale.py"],
-    srcs_version = "PY2AND3",
-    deps = [
-        "//tensorflow/python:framework",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)
-
-py_test(
-    name = "loss_scale_test",
-    size = "medium",
-    srcs = ["loss_scale_test.py"],
-    python_version = "PY2",
-    deps = [
-        ":loss_scale",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python/distribute:mirrored_strategy",
-        "//tensorflow/python/distribute:one_device_strategy",
-        "//tensorflow/python/keras",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)
-
 py_library(
     name = "test_util",
     srcs = ["test_util.py"],

tensorflow/python/keras/mixed_precision/experimental/keras_test.py

@@ -36,7 +36,6 @@
 from tensorflow.python.keras import regularizers
 from tensorflow.python.keras.engine import base_layer
 from tensorflow.python.keras.layers import core
-from tensorflow.python.keras.mixed_precision.experimental import loss_scale as loss_scale_module
 from tensorflow.python.keras.mixed_precision.experimental import loss_scale_optimizer
 from tensorflow.python.keras.mixed_precision.experimental import policy
 from tensorflow.python.keras.mixed_precision.experimental import test_util as mp_test_util
@@ -45,6 +44,7 @@
 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
 from tensorflow.python.training.tracking import util as trackable_utils
 from tensorflow.python.util import nest