Add missing build dependencies (which had been indirect dependencies previously).

Access Keras using tf.keras rather than importing keras itself.

PiperOrigin-RevId: 274047775

Author: Neural-Link Team

neural_structured_learning/BUILD

@@ -15,6 +15,7 @@
 # Description:
 # Build rules for TensorFlow Neural Structured Learning.
 
+# Placeholder for internal Python strict compatibility macro.
 # Internal annotation for sync
 
 package(default_visibility = ["//visibility:public"])
@@ -28,6 +29,7 @@ py_library(
     srcs = ["__init__.py"],
     deps = [
         ":version",
+        "//neural_structured_learning/configs",
         "//neural_structured_learning/estimator",
         "//neural_structured_learning/keras",
         "//neural_structured_learning/lib",

neural_structured_learning/configs/BUILD

@@ -15,6 +15,8 @@
 # Description:
 # Build rules for config libraries in Neural Structured Learning.
 
+# Placeholder for internal Python strict compatibility macro.
+
 package(
     default_visibility = ["//visibility:public"],
     licenses = ["notice"],  # Apache 2.0

neural_structured_learning/estimator/BUILD

@@ -15,6 +15,7 @@
 # Description:
 # Build rules for Estimator APIs in Neural Structured Learning.
 
+# Placeholder for internal Python strict compatibility macro.
 # Placeholder for internal Python version compatibility macro.
 
 package(

neural_structured_learning/keras/BUILD

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Placeholder for internal Python strict compatibility macro.
 # Placeholder for internal Python version compatibility macro.
 
 # Description:
@@ -30,6 +31,7 @@ py_library(
     deps = [
         ":adversarial_regularization",
         ":graph_regularization",
+        "//neural_structured_learning/keras/layers",
     ],
 )
 

neural_structured_learning/keras/adversarial_regularization.py

@@ -26,7 +26,6 @@
 import neural_structured_learning.lib as nsl_lib
 import six
 import tensorflow as tf
-import tensorflow.keras as keras
 
 
 def adversarial_loss(features,
@@ -193,16 +192,15 @@ def __call__(self, *args, **kwargs):
     return loss_value
 
   def _is_sparse_categorical_loss(self):
-    return self.loss_fn == keras.losses.sparse_categorical_crossentropy or (
-        isinstance(self.loss_fn, keras.losses.SparseCategoricalCrossentropy))
+    return self.loss_fn == tf.keras.losses.sparse_categorical_crossentropy or (
+        isinstance(self.loss_fn, tf.keras.losses.SparseCategoricalCrossentropy))
 
   def _is_binary_classification_loss(self):
-    return self.loss_fn in (keras.losses.binary_crossentropy,
-                            keras.losses.hinge,
-                            keras.losses.squared_hinge) or isinstance(
-                                self.loss_fn,
-                                (keras.losses.BinaryCrossentropy,
-                                 keras.losses.Hinge, keras.losses.SquaredHinge))
+    return self.loss_fn in (
+        tf.keras.losses.binary_crossentropy,
+        tf.keras.losses.hinge, tf.keras.losses.squared_hinge) or isinstance(
+            self.loss_fn, (tf.keras.losses.BinaryCrossentropy,
+                           tf.keras.losses.Hinge, tf.keras.losses.SquaredHinge))
 
   def resolve_metric(self, metric):
     """Resolves potentially ambiguous metric name based on the loss function."""
@@ -237,21 +235,21 @@ def _prepare_loss_fns(loss, output_names):
       if name not in loss:
         raise ValueError(
             'Loss for {} not found in `loss` dictionary.'.format(name))
-    return [keras.losses.get(loss[name]) for name in output_names]
+    return [tf.keras.losses.get(loss[name]) for name in output_names]
 
   # loss for single output, or shared loss fn for multiple outputs
   if isinstance(loss, six.string_types):
-    return [keras.losses.get(loss) for _ in output_names]
+    return [tf.keras.losses.get(loss) for _ in output_names]
 
   # losses for multiple outputs indexed by position
   if isinstance(loss, collections.Sequence):
     if len(loss) != len(output_names):
       raise ValueError('`loss` should have the same number of elements as '
                        'model output')
-    return six.moves.map(keras.losses.get, loss)
+    return six.moves.map(tf.keras.losses.get, loss)
 
   # loss for single output, or shared loss fn for multiple outputs
-  return [keras.losses.get(loss) for _ in output_names]
+  return [tf.keras.losses.get(loss) for _ in output_names]
 
 
 def _prepare_loss_weights(loss_weights, output_names):
@@ -294,7 +292,7 @@ def clone(metric):
     # adversarial-regularized models, and also on multiple outputs in one model.
     # The cloning logic is the same as the `clone_metric` function in
     # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/metrics.py
-    if not isinstance(metric, keras.metrics.Metric):
+    if not isinstance(metric, tf.keras.metrics.Metric):
       return metric
     with tf.init_scope():
       return metric.__class__.from_config(metric.get_config())
@@ -348,7 +346,7 @@ def _prepare_metric_fns(metrics, output_names, loss_wrappers):
   metric_fns = []
   for per_output_metrics, loss_wrapper in zip(metrics, loss_wrappers):
     metric_fns.append([
-        keras.metrics.get(loss_wrapper.resolve_metric(metric))
+        tf.keras.metrics.get(loss_wrapper.resolve_metric(metric))
         for metric in to_list(per_output_metrics)
     ])
   return metric_fns
@@ -390,15 +388,15 @@ def _compute_loss_and_metrics(losses,
       value = metric_fn(label, output)
       # Metric objects always return an aggregated result, and shouldn't be
       # aggregated again.
-      if isinstance(metric_fn, keras.metrics.Metric):
+      if isinstance(metric_fn, tf.keras.metrics.Metric):
         aggregation = None
       else:
         aggregation = 'mean'
       output_metrics.append((value, aggregation, metric_name))
   return tf.add_n(total_loss), output_metrics
 
 
-class AdversarialRegularization(keras.Model):
+class AdversarialRegularization(tf.keras.Model):
   """Wrapper thats adds adversarial regularization to a given `tf.keras.Model`.
 
   This model will reuse the layers and variables as the given `base_model`, so
@@ -567,7 +565,7 @@ def _build_labeled_metrics(self, output_names, labeled_losses):
       per_output_metrics = []
       for metric_fn in metric_fns:
         metric_name = self._make_metric_name(metric_fn, label_key)
-        if isinstance(metric_fn, keras.metrics.Metric):
+        if isinstance(metric_fn, tf.keras.metrics.Metric):
           # Updates the name of the Metric object to make sure it is unique.
           metric_fn._name = metric_name  # pylint: disable=protected-access
         per_output_metrics.append((metric_fn, metric_name))
@@ -682,7 +680,7 @@ def perturb_on_batch(self, x, **config_kwargs):
           adv_inputs,
           expand_composites=False)
     else:
-      adv_inputs = keras.backend.function([], adv_inputs)([])
+      adv_inputs = tf.keras.backend.function([], adv_inputs)([])
 
     # Inserts the labels and sample_weights back to the input dictionary, so
     # the returned input has the same structure as the original input.

neural_structured_learning/keras/layers/BUILD

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Placeholder for internal Python strict compatibility macro.
 # Placeholder for internal Python version compatibility macro.
 
 # Description:

neural_structured_learning/lib/BUILD

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Placeholder for internal Python strict compatibility macro.
 # Placeholder for internal Python version compatibility macro.
 
 # Description:

neural_structured_learning/tools/BUILD

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Placeholder for internal Python strict compatibility macro.
 # Placeholder for internal Python version compatibility macro.
 
 # Description: