Temporarily avoid isinstance(..., ProbabilisticModel)

tests/unit/models/test_interfaces.py

@@ -78,6 +78,9 @@ def _model_stack() -> (
 
 def test_model_stack_predict() -> None:
     stack, (model01, model2, model3) = _model_stack()
+    assert all(
+        isinstance(model, TrainableProbabilisticModel) for model in (stack, model01, model2, model3)
+    )
     query_points = tf.random.uniform([5, 7, 3])
     mean, var = stack.predict(query_points)
 

trieste/acquisition/rule.py

@@ -1095,7 +1095,7 @@ def with_input_active_dims(
 
         # No selection for models.
         # Nothing to do if active dimensions are not set.
-        if isinstance(value, ProbabilisticModel) or self.input_active_dims is None:
+        if not isinstance(value, (Dataset, tf.Tensor)) or self.input_active_dims is None:
             return value
 
         # Select components of query points for datasets.