Negative sampling - Fix boolean checks (#571)

* Check if None explictly to help tensorflow compile correctly

* Run negative sampling model in both eager and graph mode

* Move custom layers outside of test to avoid re-register class error

merlin/models/tf/core/prediction.py

@@ -31,9 +31,9 @@ def with_updates(
         self, targets=None, features=None, mask=None, training=None, testing=None
     ) -> "PredictionContext":
         return PredictionContext(
-            features or self.features,
-            targets or self.targets,
-            mask or self.mask,
+            features if features is not None else self.features,
+            targets if targets is not None else self.targets,
+            mask if mask is not None else self.mask,
             training or self.training,
             testing or self.testing,
         )

merlin/models/tf/models/base.py

@@ -744,8 +744,8 @@ def _call_child(
 
         outputs = call_layer(child, inputs, **call_kwargs)
         if isinstance(outputs, Prediction):
-            targets = outputs.targets or context.targets
-            features = outputs.features or context.features
+            targets = outputs.targets if outputs.targets is not None else context.targets
+            features = outputs.features if outputs.features is not None else context.features
             outputs = outputs[0]
             context = context.with_updates(targets=targets, features=features)
 

tests/unit/tf/data_augmentation/test_negative_sampling.py

@@ -27,6 +27,21 @@
 from merlin.schema import ColumnSchema, Schema, Tags
 
 
+@tf.keras.utils.register_keras_serializable(package="merlin.models")
+class ExampleIsTraining(tf.keras.layers.Layer):
+    def call(self, inputs, training=False):
+        return training
+
+
+@tf.keras.utils.register_keras_serializable(package="merlin.models")
+class ExamplePredictionIdentity(tf.keras.layers.Layer):
+    def call(self, inputs, targets=None):
+        return Prediction(inputs, targets)
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
 class TestAddRandomNegativesToBatch:
     def test_dataloader(self):
         schema = Schema(
@@ -97,27 +112,15 @@ def test_calling(self, music_streaming_data: Dataset, to_dense: bool, tf_random_
             for f in with_negatives.values()
         )
 
-    def test_in_model(self, music_streaming_data: Dataset, tf_random_seed: int):
+    @pytest.mark.parametrize("run_eagerly", [True, False])
+    def test_in_model(self, run_eagerly, music_streaming_data: Dataset, tf_random_seed: int):
         dataset = music_streaming_data
         schema = dataset.schema
 
-        @tf.keras.utils.register_keras_serializable(package="merlin.models")
-        class Training(tf.keras.layers.Layer):
-            def call(self, inputs, training=False):
-                return training
-
-        @tf.keras.utils.register_keras_serializable(package="merlin.models")
-        class PredictionIdentity(tf.keras.layers.Layer):
-            def call(self, inputs, targets=None):
-                return Prediction(inputs, targets)
-
-            def compute_output_shape(self, input_shape):
-                return input_shape
-
         sampling = mm.Cond(
-            Training(),
+            ExampleIsTraining(),
             UniformNegativeSampling(schema, 5, seed=tf_random_seed),
-            PredictionIdentity(),
+            ExamplePredictionIdentity(),
         )
         model = mm.Model(
             mm.InputBlock(schema),
@@ -137,7 +140,7 @@ def compute_output_shape(self, input_shape):
         without_negatives = model(features)
         assert without_negatives.shape[0] == batch_size
 
-        testing_utils.model_test(model, dataset)
+        testing_utils.model_test(model, dataset, run_eagerly=run_eagerly)
 
     def test_model_with_dataloader(self, music_streaming_data: Dataset, tf_random_seed: int):
         add_negatives = UniformNegativeSampling(music_streaming_data.schema, 5, seed=tf_random_seed)