Replace literal return type with None

baybe/surrogates/base.py

@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import pandas as pd
 from attrs import define, field
@@ -110,7 +110,7 @@ def to_botorch(self) -> Model:
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         """Return the scaler to be used for the given parameter."""
         return MinMaxScaler()
 
@@ -129,7 +129,7 @@ def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
         # Create the composite scaler from the parameter-wise scaler objects
         transformers = [
             (
-                self._make_parameter_scaler(p),
+                "passthrough" if (s := self._make_parameter_scaler(p) is None) else s,
                 [c for c in p.comp_rep_columns if c in searchspace.comp_rep_columns],
             )
             for p in searchspace.parameters

baybe/surrogates/gaussian_process/core.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar, Literal
+from typing import TYPE_CHECKING, ClassVar
 
 from attrs import define, field
 from attrs.validators import instance_of
@@ -115,12 +115,12 @@ def to_botorch(self) -> Model:  # noqa: D102
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         # See base class.
 
         # Task parameters are handled separately through an index kernel
         if isinstance(parameter, TaskParameter):
-            return "passthrough"
+            return
 
         return MinMaxScaler()
 

baybe/surrogates/ngboost.py

@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from attr import define, field
 from ngboost import NGBRegressor
@@ -56,11 +56,11 @@ def __attrs_post_init__(self):
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return "passthrough"
+        return
 
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:

baybe/surrogates/random_forest.py

@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 from attr import define, field
@@ -51,11 +51,11 @@ class RandomForestSurrogate(GaussianSurrogate):
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return "passthrough"
+        return
 
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]: