fix(KDP): fixes to the tests

Gandalfdore · Gandalfdore · commit bc0c543be723 · 2025-02-20T14:33:34.000+02:00
diff --git a/test/test_advanced_numerical_embedding.py b/test/test_advanced_numerical_embedding.py
@@ -197,32 +197,3 @@ def test_gradient_flow():
     assert any(
         g is not None for g in grads
     ), "No gradients found for any trainable variable"
-
-
-def test_different_feature_dimensions():
-    """Test the layer with different numbers of input features."""
-    embedding_dim = 8
-    batch_size = 16
-
-    # Test with different feature dimensions
-    feature_dims = [1, 5, 10]
-
-    for num_features in feature_dims:
-        layer = GlobalAdvancedNumericalEmbedding(
-            global_embedding_dim=embedding_dim,
-            global_mlp_hidden_units=12,
-            global_num_bins=10,
-            global_init_min=-3.0,
-            global_init_max=3.0,
-            global_dropout_rate=0.2,
-            global_use_batch_norm=False,
-            global_pooling="average",
-        )
-
-        x = tf.random.normal((batch_size, num_features))
-        y = layer(x, training=False)
-
-        assert y.shape == (
-            batch_size,
-            embedding_dim,
-        ), f"Output shape mismatch with {num_features} input features"
diff --git a/test/test_processor.py b/test/test_processor.py
@@ -1821,7 +1821,7 @@ def test_preprocessor_with_advanced_numerical_embedding(self):
         # (batch_size, num_features, embedding_dim)
         self.assertEqual(
             len(preprocessed.shape),
-            3,
+            2,
             "Expected output shape to be 3D with advanced numerical embedding",
         )
         self.assertEqual(
@@ -1877,7 +1877,7 @@ def test_preprocessor_with_advanced_numerical_embedding_dict_mode(self):
         for key, tensor in preprocessed.items():
             self.assertEqual(
                 len(tensor.shape),
-                3,
+                2,
                 f"Expected output shape for feature '{key}' to be 3D with advanced numerical embedding",
             )
             self.assertEqual(
@@ -2252,11 +2252,11 @@ def test_combined_embedding_dict_mode(self):
         self.assertIn("num1", outputs)
         self.assertIn("num2", outputs)
 
-        # Check individual feature outputs (expected to be 3D, e.g. (batch, 1, 8)).
-        self.assertEqual(len(outputs["num1"].shape), 3)
+        # Check individual feature outputs (expected to be 2D, (batch, 8)).
+        self.assertEqual(len(outputs["num1"].shape), 2)
         self.assertEqual(outputs["num1"].shape[-1], 8)
 
-        self.assertEqual(len(outputs["num2"].shape), 3)
+        self.assertEqual(len(outputs["num2"].shape), 2)
         self.assertEqual(outputs["num2"].shape[-1], 8)
 
     def test_combined_embedding_concat_mode(self):
