Fix recently added tests

secondmind-labs · Sep 4, 2024 · a37ffc7 · a37ffc7
1 parent 971f559
commit a37ffc7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tests/unit/models/test_interfaces.py b/tests/unit/models/test_interfaces.py
@@ -23,6 +23,7 @@
 import pytest
 import tensorflow as tf
 import tensorflow_probability as tfp
+from gpflow.keras import tf_keras
 
 from tests.util.misc import assert_datasets_allclose, quadratic, random_seed
 from tests.util.models.gpflow.models import (
@@ -304,7 +305,7 @@ def test_encoded_supports_predict_y() -> None:
 
 
 def test_encoded_probabilistic_model_keras_embedding() -> None:
-    encoder = tf.keras.layers.Embedding(3, 2)
+    encoder = tf_keras.layers.Embedding(3, 2)
     model = _EncodedModel(encoder=encoder)
     query_points = tf.random.uniform([3, 5], minval=0, maxval=3, dtype=tf.int32)
     mean, var = model.predict(query_points)

diff --git a/trieste/space.py b/trieste/space.py
@@ -25,6 +25,7 @@
 import tensorflow as tf
 import tensorflow_probability as tfp
 from check_shapes import check_shapes
+from gpflow.keras import tf_keras
 from typing_extensions import Protocol, runtime_checkable
 
 from .types import TensorType
@@ -643,7 +644,7 @@ def encoder(x: TensorType) -> TensorType:
                 )
             columns = tf.split(flat_x, flat_x.shape[-1], axis=1)
             encoders = [
-                tf.keras.layers.CategoryEncoding(num_tokens=len(ts), output_mode="one_hot")
+                tf_keras.layers.CategoryEncoding(num_tokens=len(ts), output_mode="one_hot")
                 for ts in self.tags
             ]
             encoded = tf.concat(