From 59e969792d5118ae13fc67f8514e09f3bff083b5 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Wed, 15 Apr 2026 11:20:16 +0200
Subject: [PATCH 01/16] Move _ParameterSelectorMixin to components/kernel.py

---
 baybe/parameters/selectors.py                 | 38 +-----------------
 .../gaussian_process/components/kernel.py     | 39 ++++++++++++++++++-
 .../gaussian_process/presets/baybe.py         |  6 ++-
 .../gaussian_process/presets/edbo.py          |  6 ++-
 .../gaussian_process/presets/edbo_smoothed.py |  6 ++-
 5 files changed, 51 insertions(+), 44 deletions(-)

diff --git a/baybe/parameters/selectors.py b/baybe/parameters/selectors.py
index 6b1a4ae16d..e72b6fe4d9 100644
--- a/baybe/parameters/selectors.py
+++ b/baybe/parameters/selectors.py
@@ -3,15 +3,13 @@
 import re
 from abc import ABC, abstractmethod
 from collections.abc import Collection
-from typing import ClassVar, Protocol
+from typing import Protocol
 
 from attrs import Converter, define, field
-from attrs.converters import optional
 from attrs.validators import deep_iterable, instance_of, min_len
 from typing_extensions import override
 
 from baybe.parameters.base import Parameter
-from baybe.searchspace.core import SearchSpace
 from baybe.utils.basic import to_tuple
 from baybe.utils.conversion import nonstring_to_tuple
 
@@ -131,37 +129,3 @@ def to_parameter_selector(
         return TypeSelector(items)
 
     raise TypeError(f"Cannot convert {x!r} to a parameter selector.")
-
-
-@define
-class _ParameterSelectorMixin:
-    """A mixin class to enable parameter selection."""
-
-    # For internal use only: sanity check mechanism to remind developers of new
-    # subclasses to actually use the parameter selector when it is provided
-    # TODO: Perhaps we can find a more elegant way to enforce this by design
-    _uses_parameter_names: ClassVar[bool] = False
-
-    parameter_selector: ParameterSelectorProtocol | None = field(
-        default=None, converter=optional(to_parameter_selector), kw_only=True
-    )
-    """An optional selector to specify which parameters are to be considered."""
-
-    def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | None:
-        """Get the names of the parameters to be considered."""
-        if self.parameter_selector is None:
-            return None
-
-        return tuple(
-            p.name for p in searchspace.parameters if self.parameter_selector(p)
-        )
-
-    def __attrs_post_init__(self):
-        if self.parameter_selector is not None and not self._uses_parameter_names:
-            raise AssertionError(
-                f"A `parameter_selector` was provided to "
-                f"`{type(self).__name__}`, but the class does not set "
-                f"`_uses_parameter_names = True`. Subclasses that accept a "
-                f"parameter selector must explicitly set this flag to confirm "
-                f"they actually use the selected parameter names."
-            )
diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index ca4da066e7..bc9b215465 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -3,9 +3,10 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
 from attrs import define, field
+from attrs.converters import optional
 from attrs.validators import is_callable
 from typing_extensions import override
 
@@ -13,7 +14,9 @@
 from baybe.kernels.composite import ProductKernel
 from baybe.parameters.categorical import TaskParameter
 from baybe.parameters.selectors import (
+    ParameterSelectorProtocol,
     TypeSelector,
+    to_parameter_selector,
 )
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.gaussian_process.components.generic import (
@@ -35,6 +38,40 @@
     PlainKernelFactory = PlainGPComponentFactory[Kernel]
 
 
+@define
+class _ParameterSelectorMixin:
+    """A mixin class to enable parameter selection."""
+
+    # For internal use only: sanity check mechanism to remind developers of new
+    # subclasses to actually use the parameter selector when it is provided
+    # TODO: Perhaps we can find a more elegant way to enforce this by design
+    _uses_parameter_names: ClassVar[bool] = False
+
+    parameter_selector: ParameterSelectorProtocol | None = field(
+        default=None, converter=optional(to_parameter_selector), kw_only=True
+    )
+    """An optional selector to specify which parameters are to be considered."""
+
+    def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | None:
+        """Get the names of the parameters to be considered."""
+        if self.parameter_selector is None:
+            return None
+
+        return tuple(
+            p.name for p in searchspace.parameters if self.parameter_selector(p)
+        )
+
+    def __attrs_post_init__(self):
+        if self.parameter_selector is not None and not self._uses_parameter_names:
+            raise AssertionError(
+                f"A `parameter_selector` was provided to "
+                f"`{type(self).__name__}`, but the class does not set "
+                f"`_uses_parameter_names = True`. Subclasses that accept a "
+                f"parameter selector must explicitly set this flag to confirm "
+                f"they actually use the selected parameter names."
+            )
+
+
 @define
 class ICMKernelFactory(KernelFactoryProtocol):
     """A kernel factory that constructs an ICM kernel for transfer learning.
diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py
index bc820ae52b..cd19242998 100644
--- a/baybe/surrogates/gaussian_process/presets/baybe.py
+++ b/baybe/surrogates/gaussian_process/presets/baybe.py
@@ -13,11 +13,13 @@
 from baybe.parameters.selectors import (
     ParameterSelectorProtocol,
     TypeSelector,
-    _ParameterSelectorMixin,
     to_parameter_selector,
 )
 from baybe.searchspace.core import SearchSpace
-from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol
+from baybe.surrogates.gaussian_process.components.kernel import (
+    KernelFactoryProtocol,
+    _ParameterSelectorMixin,
+)
 from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory
 from baybe.surrogates.gaussian_process.presets.edbo_smoothed import (
     SmoothedEDBOKernelFactory,
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 36611c258e..8204913a41 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -16,13 +16,15 @@
 from baybe.parameters.selectors import (
     ParameterSelectorProtocol,
     TypeSelector,
-    _ParameterSelectorMixin,
     to_parameter_selector,
 )
 from baybe.parameters.substance import SubstanceParameter
 from baybe.priors.basic import GammaPrior
 from baybe.searchspace.discrete import SubspaceDiscrete
-from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol
+from baybe.surrogates.gaussian_process.components.kernel import (
+    KernelFactoryProtocol,
+    _ParameterSelectorMixin,
+)
 from baybe.surrogates.gaussian_process.components.likelihood import (
     LikelihoodFactoryProtocol,
 )
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index eb8d57491a..77fc74354b 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -15,11 +15,13 @@
 from baybe.parameters.selectors import (
     ParameterSelectorProtocol,
     TypeSelector,
-    _ParameterSelectorMixin,
     to_parameter_selector,
 )
 from baybe.priors.basic import GammaPrior
-from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol
+from baybe.surrogates.gaussian_process.components.kernel import (
+    KernelFactoryProtocol,
+    _ParameterSelectorMixin,
+)
 from baybe.surrogates.gaussian_process.components.likelihood import (
     LikelihoodFactoryProtocol,
 )

From bcc51bbb04e8f80220db64f5f63e28e85b2b157c Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Wed, 15 Apr 2026 13:23:08 +0200
Subject: [PATCH 02/16] Turn mixin class back into ABC

---
 .../gaussian_process/components/kernel.py     | 20 +++++++++++++------
 .../gaussian_process/presets/baybe.py         |  4 ++--
 .../gaussian_process/presets/edbo.py          |  5 ++---
 .../gaussian_process/presets/edbo_smoothed.py |  5 ++---
 docs/conf.py                                  |  2 +-
 5 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index bc9b215465..5d9aecfc35 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from functools import partial
 from typing import TYPE_CHECKING, ClassVar
 
@@ -39,21 +40,21 @@
 
 
 @define
-class _ParameterSelectorMixin:
-    """A mixin class to enable parameter selection."""
+class _KernelFactory(KernelFactoryProtocol, ABC):
+    """Base class for kernel factories."""
 
     # For internal use only: sanity check mechanism to remind developers of new
-    # subclasses to actually use the parameter selector when it is provided
+    # factories to actually use the parameter selector when it is provided
     # TODO: Perhaps we can find a more elegant way to enforce this by design
     _uses_parameter_names: ClassVar[bool] = False
 
     parameter_selector: ParameterSelectorProtocol | None = field(
-        default=None, converter=optional(to_parameter_selector), kw_only=True
+        default=None, converter=optional(to_parameter_selector)
     )
-    """An optional selector to specify which parameters are to be considered."""
+    """An optional selector to specify which parameters are considered by the kernel."""
 
     def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | None:
-        """Get the names of the parameters to be considered."""
+        """Get the names of the parameters to be considered by the kernel."""
         if self.parameter_selector is None:
             return None
 
@@ -71,6 +72,13 @@ def __attrs_post_init__(self):
                 f"they actually use the selected parameter names."
             )
 
+    @override
+    @abstractmethod
+    def __call__(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> Kernel | GPyTorchKernel:
+        pass
+
 
 @define
 class ICMKernelFactory(KernelFactoryProtocol):
diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py
index cd19242998..be20f5cd19 100644
--- a/baybe/surrogates/gaussian_process/presets/baybe.py
+++ b/baybe/surrogates/gaussian_process/presets/baybe.py
@@ -18,7 +18,7 @@
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.gaussian_process.components.kernel import (
     KernelFactoryProtocol,
-    _ParameterSelectorMixin,
+    _KernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory
 from baybe.surrogates.gaussian_process.presets.edbo_smoothed import (
@@ -50,7 +50,7 @@ def __call__(
 
 
 @define
-class BayBETaskKernelFactory(KernelFactoryProtocol, _ParameterSelectorMixin):
+class BayBETaskKernelFactory(_KernelFactory):
     """The factory providing the default task kernel for Gaussian process surrogates."""
 
     _uses_parameter_names: ClassVar[bool] = True
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 8204913a41..6e0dd06eeb 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -22,8 +22,7 @@
 from baybe.priors.basic import GammaPrior
 from baybe.searchspace.discrete import SubspaceDiscrete
 from baybe.surrogates.gaussian_process.components.kernel import (
-    KernelFactoryProtocol,
-    _ParameterSelectorMixin,
+    _KernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
     LikelihoodFactoryProtocol,
@@ -58,7 +57,7 @@ def _contains_encoding(
 
 
 @define
-class EDBOKernelFactory(KernelFactoryProtocol, _ParameterSelectorMixin):
+class EDBOKernelFactory(_KernelFactory):
     """A factory providing EDBO kernels.
 
     References:
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index 77fc74354b..eb7eca36da 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -19,8 +19,7 @@
 )
 from baybe.priors.basic import GammaPrior
 from baybe.surrogates.gaussian_process.components.kernel import (
-    KernelFactoryProtocol,
-    _ParameterSelectorMixin,
+    _KernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
     LikelihoodFactoryProtocol,
@@ -40,7 +39,7 @@
 
 
 @define
-class SmoothedEDBOKernelFactory(KernelFactoryProtocol, _ParameterSelectorMixin):
+class SmoothedEDBOKernelFactory(_KernelFactory):
     """A factory providing smoothed versions of EDBO kernels.
 
     Takes the low and high dimensional limits of
diff --git a/docs/conf.py b/docs/conf.py
index 5251bb8654..6096369714 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -150,7 +150,7 @@
     ("py:class", "baybe.acquisition.acqfs._ExpectedHypervolumeImprovement"),
     ("py:class", "baybe.settings._SlottedContextDecorator"),
     ("py:class", "baybe.surrogates.gaussian_process.components.PlainKernelFactory"),
-    ("py:class", "baybe.parameters.selectors._ParameterSelectorMixin"),
+    ("py:class", "baybe.surrogates.gaussian_process.components.kernel._KernelFactory"),
     # Deprecation
     ("py:.*", "baybe.targets._deprecated.*"),
 ]

From 9c6051c771e5935b2f94320bc5bdbdfad8ef7671 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 2 Apr 2026 09:45:50 +0200
Subject: [PATCH 03/16] Add ParameterKind Flag enum and Parameter.kind property

---
 baybe/parameters/__init__.py |  2 ++
 baybe/parameters/base.py     |  8 ++++++++
 baybe/parameters/enum.py     | 34 +++++++++++++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/baybe/parameters/__init__.py b/baybe/parameters/__init__.py
index 93e62b6ee9..6fe2c0d8f4 100644
--- a/baybe/parameters/__init__.py
+++ b/baybe/parameters/__init__.py
@@ -5,6 +5,7 @@
 from baybe.parameters.enum import (
     CategoricalEncoding,
     CustomEncoding,
+    ParameterKind,
     SubstanceEncoding,
 )
 from baybe.parameters.numerical import (
@@ -22,6 +23,7 @@
     "MeasurableMetadata",
     "NumericalContinuousParameter",
     "NumericalDiscreteParameter",
+    "ParameterKind",
     "SubstanceEncoding",
     "SubstanceParameter",
     "TaskParameter",
diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py
index 2d4df2bc77..7986a21124 100644
--- a/baybe/parameters/base.py
+++ b/baybe/parameters/base.py
@@ -21,6 +21,7 @@
 from baybe.utils.metadata import MeasurableMetadata, to_metadata
 
 if TYPE_CHECKING:
+    from baybe.parameters.enum import ParameterKind
     from baybe.searchspace.continuous import SubspaceContinuous
     from baybe.searchspace.core import SearchSpace
     from baybe.searchspace.discrete import SubspaceDiscrete
@@ -77,6 +78,13 @@ def is_discrete(self) -> bool:
         """Boolean indicating if this is a discrete parameter."""
         return isinstance(self, DiscreteParameter)
 
+    @property
+    def kind(self) -> ParameterKind:
+        """The kind of the parameter."""
+        from baybe.parameters.enum import ParameterKind
+
+        return ParameterKind.from_parameter(self)
+
     @property
     @abstractmethod
     def comp_rep_columns(self) -> tuple[str, ...]:
diff --git a/baybe/parameters/enum.py b/baybe/parameters/enum.py
index 3161f67dc7..07e9f9fb34 100644
--- a/baybe/parameters/enum.py
+++ b/baybe/parameters/enum.py
@@ -1,6 +1,38 @@
 """Parameter-related enumerations."""
 
-from enum import Enum
+from __future__ import annotations
+
+from enum import Enum, Flag, auto
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from baybe.parameters.base import Parameter
+
+
+class ParameterKind(Flag):
+    """Flag enum encoding the kind of a parameter.
+
+    Can be used to express compatibility (e.g. Gaussian process kernel factories)
+    with different parameter types via bitwise combination of flags.
+    """
+
+    REGULAR = auto()
+    """Regular parameter undergoing no special treatment."""
+
+    TASK = auto()
+    """Task parameter for transfer learning."""
+
+    FIDELITY = auto()
+    """Fidelity parameter for multi-fidelity modelling."""
+
+    @staticmethod
+    def from_parameter(parameter: Parameter) -> ParameterKind:
+        """Determine the kind of a parameter from its type."""
+        from baybe.parameters.categorical import TaskParameter
+
+        if isinstance(parameter, TaskParameter):
+            return ParameterKind.TASK
+        return ParameterKind.REGULAR
 
 
 class ParameterEncoding(Enum):

From 447cf6afd10ce3f4aee221239a88946b11913d31 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Wed, 15 Apr 2026 13:26:51 +0200
Subject: [PATCH 04/16] Wire parameter kind validation into kernel factories
 via template method

---
 .../gaussian_process/components/kernel.py     | 52 ++++++++++++++++---
 .../gaussian_process/presets/baybe.py         | 20 ++++---
 .../gaussian_process/presets/edbo.py          |  2 +-
 .../gaussian_process/presets/edbo_smoothed.py |  2 +-
 4 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index 5d9aecfc35..350a8405bf 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from functools import partial
 from typing import TYPE_CHECKING, ClassVar
 
@@ -11,9 +12,11 @@
 from attrs.validators import is_callable
 from typing_extensions import override
 
+from baybe.exceptions import IncompatibleSearchSpaceError
 from baybe.kernels.base import Kernel
 from baybe.kernels.composite import ProductKernel
 from baybe.parameters.categorical import TaskParameter
+from baybe.parameters.enum import ParameterKind
 from baybe.parameters.selectors import (
     ParameterSelectorProtocol,
     TypeSelector,
@@ -31,6 +34,8 @@
     from gpytorch.kernels import Kernel as GPyTorchKernel
     from torch import Tensor
 
+    from baybe.parameters.base import Parameter
+
     KernelFactoryProtocol = GPComponentFactoryProtocol[Kernel | GPyTorchKernel]
     PlainKernelFactory = PlainGPComponentFactory[Kernel | GPyTorchKernel]
 else:
@@ -48,6 +53,9 @@ class _KernelFactory(KernelFactoryProtocol, ABC):
     # TODO: Perhaps we can find a more elegant way to enforce this by design
     _uses_parameter_names: ClassVar[bool] = False
 
+    supported_parameter_kinds: ClassVar[ParameterKind] = ParameterKind.REGULAR
+    """The parameter kinds supported by the kernel factory."""
+
     parameter_selector: ParameterSelectorProtocol | None = field(
         default=None, converter=optional(to_parameter_selector)
     )
@@ -62,6 +70,43 @@ def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | Non
             p.name for p in searchspace.parameters if self.parameter_selector(p)
         )
 
+    def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None:
+        """Validate that the given parameters are supported by the factory.
+
+        Args:
+            parameters: The parameters to validate.
+
+        Raises:
+            IncompatibleSearchSpaceError: If unsupported parameter kinds are found.
+        """
+        if unsupported := [
+            p.name for p in parameters if not (p.kind & self.supported_parameter_kinds)
+        ]:
+            raise IncompatibleSearchSpaceError(
+                f"'{type(self).__name__}' does not support parameter kind(s) for "
+                f"parameter(s) {unsupported}. Supported kinds: "
+                f"{self.supported_parameter_kinds}."
+            )
+
+    @override
+    def __call__(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> Kernel:
+        """Construct the kernel, validating parameter kinds before construction."""
+        if self.parameter_selector is not None:
+            params = [p for p in searchspace.parameters if self.parameter_selector(p)]
+        else:
+            params = list(searchspace.parameters)
+        self._validate_parameter_kinds(params)
+
+        return self._make(searchspace, train_x, train_y)
+
+    @abstractmethod
+    def _make(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> Kernel:
+        """Construct the kernel."""
+
     def __attrs_post_init__(self):
         if self.parameter_selector is not None and not self._uses_parameter_names:
             raise AssertionError(
@@ -72,13 +117,6 @@ def __attrs_post_init__(self):
                 f"they actually use the selected parameter names."
             )
 
-    @override
-    @abstractmethod
-    def __call__(
-        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel | GPyTorchKernel:
-        pass
-
 
 @define
 class ICMKernelFactory(KernelFactoryProtocol):
diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py
index be20f5cd19..e96a08e6cc 100644
--- a/baybe/surrogates/gaussian_process/presets/baybe.py
+++ b/baybe/surrogates/gaussian_process/presets/baybe.py
@@ -10,16 +10,14 @@
 from baybe.kernels.base import Kernel
 from baybe.kernels.basic import IndexKernel
 from baybe.parameters.categorical import TaskParameter
+from baybe.parameters.enum import ParameterKind
 from baybe.parameters.selectors import (
     ParameterSelectorProtocol,
     TypeSelector,
     to_parameter_selector,
 )
 from baybe.searchspace.core import SearchSpace
-from baybe.surrogates.gaussian_process.components.kernel import (
-    KernelFactoryProtocol,
-    _KernelFactory,
-)
+from baybe.surrogates.gaussian_process.components.kernel import _KernelFactory
 from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory
 from baybe.surrogates.gaussian_process.presets.edbo_smoothed import (
     SmoothedEDBOKernelFactory,
@@ -31,11 +29,16 @@
 
 
 @define
-class BayBEKernelFactory(KernelFactoryProtocol):
+class BayBEKernelFactory(_KernelFactory):
     """The default kernel factory for Gaussian process surrogates."""
 
+    supported_parameter_kinds: ClassVar[ParameterKind] = (
+        ParameterKind.REGULAR | ParameterKind.TASK
+    )
+    # See base class.
+
     @override
-    def __call__(
+    def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
         from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory
@@ -56,6 +59,9 @@ class BayBETaskKernelFactory(_KernelFactory):
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
 
+    supported_parameter_kinds: ClassVar[ParameterKind] = ParameterKind.TASK
+    # See base class.
+
     parameter_selector: ParameterSelectorProtocol | None = field(
         factory=lambda: TypeSelector([TaskParameter]),
         converter=to_parameter_selector,
@@ -63,7 +69,7 @@ class BayBETaskKernelFactory(_KernelFactory):
     # TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
 
     @override
-    def __call__(
+    def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
         return IndexKernel(
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 6e0dd06eeb..758c643770 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -75,7 +75,7 @@ class EDBOKernelFactory(_KernelFactory):
     # TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
 
     @override
-    def __call__(
+    def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
         effective_dims = train_x.shape[-1] - len(
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index eb7eca36da..ad5bcf7f53 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -57,7 +57,7 @@ class SmoothedEDBOKernelFactory(_KernelFactory):
     # TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
 
     @override
-    def __call__(
+    def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
         effective_dims = train_x.shape[-1] - len(

From 5845f0fbf74c1963da00e77859accf5ca9ca664b Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 2 Apr 2026 13:37:43 +0200
Subject: [PATCH 05/16] Remove dead TaskParameter filtering in kernel factories

The new parameter kind validation step guards us
---
 baybe/surrogates/gaussian_process/presets/edbo.py          | 4 +---
 baybe/surrogates/gaussian_process/presets/edbo_smoothed.py | 4 +---
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 758c643770..ecf24de106 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -78,9 +78,7 @@ class EDBOKernelFactory(_KernelFactory):
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
-        effective_dims = train_x.shape[-1] - len(
-            [p for p in searchspace.parameters if isinstance(p, TaskParameter)]
-        )
+        effective_dims = train_x.shape[-1]
 
         switching_condition = _contains_encoding(
             searchspace.discrete, _EDBO_ENCODINGS
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index ad5bcf7f53..89cf9a9ef2 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -60,9 +60,7 @@ class SmoothedEDBOKernelFactory(_KernelFactory):
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
-        effective_dims = train_x.shape[-1] - len(
-            [p for p in searchspace.parameters if isinstance(p, TaskParameter)]
-        )
+        effective_dims = train_x.shape[-1]
 
         # Interpolate prior moments linearly between low D and high D regime.
         # The high D regime itself is the average of the EDBO OHE and Mordred regime.

From 741997b9ffda88ecbbdaeecbe8f1acc5157e2511 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 2 Apr 2026 13:46:28 +0200
Subject: [PATCH 06/16] Update CHANGELOG.md

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 680f11da28..c7715778a3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate`
 - `TypeSelector` and `NameSelector` classes for parameter selection in kernel factories
 - `parameter_names` attribute to basic kernels for controlling the considered parameters
+- `ParameterKind` flag enum for classifying parameters by their role and automatic
+  parameter kind validation in kernel factories
 - `IndexKernel` and `PositiveIndexKernel` classes
 - Interpoint constraints for continuous search spaces
 - `IndexKernel` and `PositiveIndexKernel` classes

From 762b8b90450d60a6a4b973a6402b7d02874ebcb0 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 16 Apr 2026 09:08:33 +0200
Subject: [PATCH 07/16] Add parameter support test

---
 tests/test_kernel_factories.py | 75 ++++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)
 create mode 100644 tests/test_kernel_factories.py

diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py
new file mode 100644
index 0000000000..6a8acda6a6
--- /dev/null
+++ b/tests/test_kernel_factories.py
@@ -0,0 +1,75 @@
+"""Tests for kernel factories."""
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+from pytest import param
+
+from baybe.exceptions import IncompatibleSearchSpaceError
+from baybe.parameters.categorical import CategoricalParameter, TaskParameter
+from baybe.parameters.numerical import (
+    NumericalContinuousParameter,
+    NumericalDiscreteParameter,
+)
+from baybe.searchspace.core import SearchSpace
+from baybe.surrogates.gaussian_process.presets.baybe import (
+    BayBEKernelFactory,
+    BayBENumericalKernelFactory,
+    BayBETaskKernelFactory,
+)
+
+# A selector that accepts all parameters
+_SELECT_ALL = lambda parameter: True  # noqa: E731
+
+
+@pytest.mark.parametrize(
+    ("factory", "parameters", "error"),
+    [
+        param(
+            BayBENumericalKernelFactory(parameter_selector=_SELECT_ALL),
+            [TaskParameter("task", ["t1", "t2"])],
+            IncompatibleSearchSpaceError,
+            id="regular_rejects_task",
+        ),
+        param(
+            BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
+            [CategoricalParameter("cat", ["a", "b"])],
+            IncompatibleSearchSpaceError,
+            id="task_rejects_categorical",
+        ),
+        param(
+            BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
+            [NumericalDiscreteParameter("num", [1, 2, 3])],
+            IncompatibleSearchSpaceError,
+            id="task_rejects_numerical_discrete",
+        ),
+        param(
+            BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
+            [NumericalContinuousParameter("cont", (0, 1))],
+            IncompatibleSearchSpaceError,
+            id="task_rejects_numerical_continuous",
+        ),
+        param(
+            BayBEKernelFactory(),
+            [
+                NumericalContinuousParameter("cont", (0, 1)),
+                TaskParameter("task", ["t1", "t2"]),
+            ],
+            None,
+            id="combined_accepts_both",
+        ),
+    ],
+)
+def test_factory_parameter_kind_validation(factory, parameters, error):
+    """Factories reject unsupported parameter kinds and accept supported ones."""
+    ss = SearchSpace.from_product(parameters)
+    train_x = torch.zeros(2, len(ss.comp_rep_columns))
+    train_y = torch.zeros(2, 1)
+
+    with (
+        nullcontext()
+        if error is None
+        else pytest.raises(error, match="does not support")
+    ):
+        factory(ss, train_x, train_y)

From d6f3128998031059f0532d30faa2f4cb8af2ab38 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 16 Apr 2026 10:11:26 +0200
Subject: [PATCH 08/16] Make ICMKernelFactory gpytorch-compatible

---
 baybe/surrogates/gaussian_process/components/kernel.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index 350a8405bf..65b6e5e454 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -14,7 +14,6 @@
 
 from baybe.exceptions import IncompatibleSearchSpaceError
 from baybe.kernels.base import Kernel
-from baybe.kernels.composite import ProductKernel
 from baybe.parameters.categorical import TaskParameter
 from baybe.parameters.enum import ParameterKind
 from baybe.parameters.selectors import (
@@ -161,4 +160,8 @@ def __call__(
     ) -> Kernel:
         base_kernel = self.base_kernel_factory(searchspace, train_x, train_y)
         task_kernel = self.task_kernel_factory(searchspace, train_x, train_y)
-        return ProductKernel([base_kernel, task_kernel])
+        if isinstance(base_kernel, Kernel):
+            base_kernel = base_kernel.to_gpytorch(searchspace)
+        if isinstance(task_kernel, Kernel):
+            task_kernel = task_kernel.to_gpytorch(searchspace)
+        return base_kernel * task_kernel

From aaee9fcf1f69f2e1f7560d66202b67e0198b01bc Mon Sep 17 00:00:00 2001
From: emiliencgm <emilien.chen@chimieparistech.psl.eu>
Date: Thu, 26 Feb 2026 16:45:06 +0100
Subject: [PATCH 09/16] Add CHENKernelFactory draft

---
 .../gaussian_process/presets/chen.py          | 62 +++++++++++++++++++
 1 file changed, 62 insertions(+)
 create mode 100644 baybe/surrogates/gaussian_process/presets/chen.py

diff --git a/baybe/surrogates/gaussian_process/presets/chen.py b/baybe/surrogates/gaussian_process/presets/chen.py
new file mode 100644
index 0000000000..12cb9ff03e
--- /dev/null
+++ b/baybe/surrogates/gaussian_process/presets/chen.py
@@ -0,0 +1,62 @@
+"""Adaptive Prior proposed in the paper: 
+Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
+DOI: https://doi.org/10.26434/chemrxiv.10001986/v2"""
+
+from __future__ import annotations
+
+import gc
+from typing import TYPE_CHECKING
+
+import numpy as np
+from attrs import define
+from typing_extensions import override
+
+from baybe.kernels.basic import MaternKernel
+from baybe.kernels.composite import ScaleKernel
+from baybe.parameters import TaskParameter
+from baybe.priors.basic import GammaPrior
+from baybe.surrogates.gaussian_process.kernel_factory import KernelFactory
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+    from baybe.kernels.base import Kernel
+    from baybe.searchspace.core import SearchSpace
+
+import math
+
+@define
+class CHENKernelFactory(KernelFactory):
+    """ Surrogate model with an adaptive hyperprior proposed in the paper: 
+        Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
+        DOI: https://doi.org/10.26434/chemrxiv.10001986/v2
+    """
+
+    @override
+    def __call__(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> Kernel:
+        effective_dims = train_x.shape[-1] - len(
+            [p for p in searchspace.parameters if isinstance(p, TaskParameter)]
+        )
+
+        x = math.sqrt(effective_dims)
+        l_mean = 0.4 * x + 4.0
+        
+        lengthscale_prior = GammaPrior(2.0*l_mean, 2.0)
+        lengthscale_initial_value = l_mean
+        outputscale_prior = GammaPrior(1.0*l_mean, 1.0) 
+        outputscale_initial_value = l_mean
+
+        return ScaleKernel(
+            MaternKernel(
+                nu=2.5,
+                lengthscale_prior=lengthscale_prior,
+                lengthscale_initial_value=lengthscale_initial_value,
+            ),
+            outputscale_prior=outputscale_prior,
+            outputscale_initial_value=outputscale_initial_value,
+        )
+
+# Collect leftover original slotted classes processed by `attrs.define`
+gc.collect()

From 923f3465ceac1d90ffa8a1baec95ec2ffae5df8c Mon Sep 17 00:00:00 2001
From: maxfleck <maxi_fleck@posteo.com>
Date: Thu, 26 Feb 2026 21:33:36 +0100
Subject: [PATCH 10/16] Import CHENKernelFactory in __init__.py

---
 baybe/surrogates/gaussian_process/presets/__init__.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/baybe/surrogates/gaussian_process/presets/__init__.py b/baybe/surrogates/gaussian_process/presets/__init__.py
index deb7de9e64..434fbf560f 100644
--- a/baybe/surrogates/gaussian_process/presets/__init__.py
+++ b/baybe/surrogates/gaussian_process/presets/__init__.py
@@ -7,6 +7,9 @@
     BayBEMeanFactory,
 )
 
+# Chen preset
+from baybe.surrogates.gaussian_process.presets.chen import CHENKernelFactory
+
 # Core
 from baybe.surrogates.gaussian_process.presets.core import GaussianProcessPreset
 
@@ -31,6 +34,8 @@
     "BayBEKernelFactory",
     "BayBELikelihoodFactory",
     "BayBEMeanFactory",
+    # Chen preset
+    "CHENKernelFactory",
     # EDBO preset
     "EDBOKernelFactory",
     "EDBOLikelihoodFactory",

From 23ed1debfb1b969f668b4cbe7eb64ada04e2fc73 Mon Sep 17 00:00:00 2001
From: tstuyver <Thijs.stuyver@chimieparistech.psl.eu>
Date: Fri, 27 Feb 2026 00:17:44 +0100
Subject: [PATCH 11/16] Update CONTRIBUTORS.md

---
 CONTRIBUTORS.md | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
index 6270ff77c0..f639c03a8a 100644
--- a/CONTRIBUTORS.md
+++ b/CONTRIBUTORS.md
@@ -39,4 +39,10 @@
 - Kathrin Skubch (Merck KGaA, Darmstadt, Germany):\
   Transfer learning regression benchmarks infrastructure
 - Myra Zmarsly (Merck Life Science KGaA, Darmstadt, Germany):\
-  Identification of non-dominated parameter configurations
\ No newline at end of file
+  Identification of non-dominated parameter configurations
+- Thijs Stuyver (PSL University, Paris, France):\  
+  Adaptive hyper-prior tailored for reaction yield optimization tasks
+- Maximilian Fleck (PSL University, Paris, France):\
+  Adaptive hyper-prior tailored for reaction yield optimization tasks
+- Guanming Chen (PSL University, Paris, France):\
+  Adaptive hyper-prior tailored for reaction yield optimization tasks

From 05c2989185a790fdbb2d39e7c23e264f43ef29d8 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 16 Apr 2026 10:33:06 +0200
Subject: [PATCH 12/16] Apply ruff formatting

---
 .../gaussian_process/presets/chen.py          | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/presets/chen.py b/baybe/surrogates/gaussian_process/presets/chen.py
index 12cb9ff03e..069fbc6616 100644
--- a/baybe/surrogates/gaussian_process/presets/chen.py
+++ b/baybe/surrogates/gaussian_process/presets/chen.py
@@ -1,13 +1,13 @@
-"""Adaptive Prior proposed in the paper: 
+"""Adaptive Prior proposed in the paper:
 Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
-DOI: https://doi.org/10.26434/chemrxiv.10001986/v2"""
+DOI: https://doi.org/10.26434/chemrxiv.10001986/v2
+"""  # noqa
 
 from __future__ import annotations
 
 import gc
 from typing import TYPE_CHECKING
 
-import numpy as np
 from attrs import define
 from typing_extensions import override
 
@@ -25,12 +25,13 @@
 
 import math
 
+
 @define
 class CHENKernelFactory(KernelFactory):
-    """ Surrogate model with an adaptive hyperprior proposed in the paper: 
-        Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
-        DOI: https://doi.org/10.26434/chemrxiv.10001986/v2
-    """
+    """Surrogate model with an adaptive hyperprior proposed in the paper:
+    Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
+    DOI: https://doi.org/10.26434/chemrxiv.10001986/v2
+    """  # noqa
 
     @override
     def __call__(
@@ -42,10 +43,10 @@ def __call__(
 
         x = math.sqrt(effective_dims)
         l_mean = 0.4 * x + 4.0
-        
-        lengthscale_prior = GammaPrior(2.0*l_mean, 2.0)
+
+        lengthscale_prior = GammaPrior(2.0 * l_mean, 2.0)
         lengthscale_initial_value = l_mean
-        outputscale_prior = GammaPrior(1.0*l_mean, 1.0) 
+        outputscale_prior = GammaPrior(1.0 * l_mean, 1.0)
         outputscale_initial_value = l_mean
 
         return ScaleKernel(
@@ -58,5 +59,6 @@ def __call__(
             outputscale_initial_value=outputscale_initial_value,
         )
 
+
 # Collect leftover original slotted classes processed by `attrs.define`
 gc.collect()

From 5aeacdf3eeb3380aa59bb6b1c7d6552dd88ace67 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 16 Apr 2026 10:46:17 +0200
Subject: [PATCH 13/16] Finalize draft

---
 .../gaussian_process/presets/chen.py          | 41 +++++++++++--------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/presets/chen.py b/baybe/surrogates/gaussian_process/presets/chen.py
index 069fbc6616..7937ae7f6f 100644
--- a/baybe/surrogates/gaussian_process/presets/chen.py
+++ b/baybe/surrogates/gaussian_process/presets/chen.py
@@ -6,16 +6,20 @@
 from __future__ import annotations
 
 import gc
-from typing import TYPE_CHECKING
+import math
+from typing import TYPE_CHECKING, ClassVar
 
 from attrs import define
 from typing_extensions import override
 
 from baybe.kernels.basic import MaternKernel
 from baybe.kernels.composite import ScaleKernel
-from baybe.parameters import TaskParameter
 from baybe.priors.basic import GammaPrior
-from baybe.surrogates.gaussian_process.kernel_factory import KernelFactory
+from baybe.surrogates.gaussian_process.components.kernel import _KernelFactory
+from baybe.surrogates.gaussian_process.presets.baybe import (
+    BayBELikelihoodFactory,
+    BayBEMeanFactory,
+)
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -23,37 +27,33 @@
     from baybe.kernels.base import Kernel
     from baybe.searchspace.core import SearchSpace
 
-import math
-
 
 @define
-class CHENKernelFactory(KernelFactory):
+class CHENKernelFactory(_KernelFactory):
     """Surrogate model with an adaptive hyperprior proposed in the paper:
     Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
     DOI: https://doi.org/10.26434/chemrxiv.10001986/v2
     """  # noqa
 
+    _uses_parameter_names: ClassVar[bool] = True
+    # See base class.
+
     @override
-    def __call__(
+    def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
-        effective_dims = train_x.shape[-1] - len(
-            [p for p in searchspace.parameters if isinstance(p, TaskParameter)]
-        )
-
-        x = math.sqrt(effective_dims)
-        l_mean = 0.4 * x + 4.0
-
-        lengthscale_prior = GammaPrior(2.0 * l_mean, 2.0)
-        lengthscale_initial_value = l_mean
-        outputscale_prior = GammaPrior(1.0 * l_mean, 1.0)
-        outputscale_initial_value = l_mean
+        lengthscale = 0.4 * math.sqrt(train_x.shape[-1]) + 4.0
+        lengthscale_prior = GammaPrior(2.0 * lengthscale, 2.0)
+        lengthscale_initial_value = lengthscale
+        outputscale_prior = GammaPrior(1.0 * lengthscale, 1.0)
+        outputscale_initial_value = lengthscale
 
         return ScaleKernel(
             MaternKernel(
                 nu=2.5,
                 lengthscale_prior=lengthscale_prior,
                 lengthscale_initial_value=lengthscale_initial_value,
+                parameter_names=self.get_parameter_names(searchspace),
             ),
             outputscale_prior=outputscale_prior,
             outputscale_initial_value=outputscale_initial_value,
@@ -62,3 +62,8 @@ def __call__(
 
 # Collect leftover original slotted classes processed by `attrs.define`
 gc.collect()
+
+# Aliases for generic preset imports
+PresetKernelFactory = CHENKernelFactory
+PresetMeanFactory = BayBEMeanFactory
+PresetLikelihoodFactory = BayBELikelihoodFactory

From 042bc3b169aa5532d199387c8cec64f1005d85e0 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 16 Apr 2026 11:36:17 +0200
Subject: [PATCH 14/16] Add proper paper references

---
 .../gaussian_process/presets/chen.py          | 10 ++--------
 .../gaussian_process/presets/edbo.py          | 14 +++++---------
 .../gaussian_process/presets/edbo_smoothed.py | 10 +++++-----
 docs/index.md                                 |  1 +
 docs/misc/references.md                       |  4 ++++
 docs/references.bib                           | 19 +++++++++++++++++++
 docs/userguide/transfer_learning.md           |  3 ---
 7 files changed, 36 insertions(+), 25 deletions(-)
 create mode 100644 docs/misc/references.md

diff --git a/baybe/surrogates/gaussian_process/presets/chen.py b/baybe/surrogates/gaussian_process/presets/chen.py
index 7937ae7f6f..014cdb15b1 100644
--- a/baybe/surrogates/gaussian_process/presets/chen.py
+++ b/baybe/surrogates/gaussian_process/presets/chen.py
@@ -1,7 +1,4 @@
-"""Adaptive Prior proposed in the paper:
-Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
-DOI: https://doi.org/10.26434/chemrxiv.10001986/v2
-"""  # noqa
+"""Preset for adaptive kernel hyperpriors proposed by :cite:p:`Chen2026`."""
 
 from __future__ import annotations
 
@@ -30,10 +27,7 @@
 
 @define
 class CHENKernelFactory(_KernelFactory):
-    """Surrogate model with an adaptive hyperprior proposed in the paper:
-    Guanming Chen, Maximilian Fleck, Thijs Stuyver. Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors. ChemRxiv. 09 February 2026.
-    DOI: https://doi.org/10.26434/chemrxiv.10001986/v2
-    """  # noqa
+    """A factory providing adaptive hyperprior kernels as proposed by :cite:p:`Chen2026`."""  # noqa: E501
 
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index ecf24de106..fac4480e73 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -1,4 +1,4 @@
-"""EDBO preset for Gaussian process surrogates."""
+"""EDBO preset :cite:p:`Shields2021`."""
 
 from __future__ import annotations
 
@@ -58,11 +58,9 @@ def _contains_encoding(
 
 @define
 class EDBOKernelFactory(_KernelFactory):
-    """A factory providing EDBO kernels.
+    """A factory providing EDBO kernels, as proposed by :cite:p:`Shields2021`.
 
-    References:
-        * https://github.com/b-shields/edbo/blob/master/edbo/bro.py#L664
-        * https://doi.org/10.1038/s41586-021-03213-y
+    Github repository: https://github.com/b-shields/edbo
     """
 
     _uses_parameter_names: ClassVar[bool] = True
@@ -130,11 +128,9 @@ def _make(
 
 @define
 class EDBOLikelihoodFactory(LikelihoodFactoryProtocol):
-    """A factory providing EDBO likelihoods.
+    """A factory providing EDBO likelihoods, as proposed by :cite:p:`Shields2021`.
 
-    References:
-        * https://github.com/b-shields/edbo/blob/master/edbo/bro.py#L664
-        * https://doi.org/10.1038/s41586-021-03213-y
+    Github repository: https://github.com/b-shields/edbo
     """
 
     @override
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index 89cf9a9ef2..16b3c2c26c 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -1,4 +1,4 @@
-"""Smoothed EDBO preset for Gaussian process surrogates."""
+"""Smoothed EDBO preset (adapted from :cite:p:`Shields2021`)."""
 
 from __future__ import annotations
 
@@ -40,12 +40,12 @@
 
 @define
 class SmoothedEDBOKernelFactory(_KernelFactory):
-    """A factory providing smoothed versions of EDBO kernels.
+    """A factory providing smoothed versions of EDBO kernels (adapted from :cite:p:`Shields2021`).
 
     Takes the low and high dimensional limits of
     :class:`baybe.surrogates.gaussian_process.presets.edbo.EDBOKernelFactory`
     and interpolates the prior moments linearly in between.
-    """
+    """  # noqa: E501
 
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
@@ -94,12 +94,12 @@ def _make(
 
 @define
 class SmoothedEDBOLikelihoodFactory(LikelihoodFactoryProtocol):
-    """A factory providing smoothed versions of EDBO likelihoods.
+    """A factory providing smoothed versions of EDBO likelihoods (adapted from :cite:p:`Shields2021`).
 
     Takes the low and high dimensional limits of
     :class:`baybe.surrogates.gaussian_process.presets.edbo.EDBOLikelihoodFactory`
     and interpolates the prior moments linearly in between.
-    """
+    """  # noqa: E501
 
     @override
     def __call__(
diff --git a/docs/index.md b/docs/index.md
index 3f8ba25cd9..31d138a369 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -40,6 +40,7 @@ Contribute <misc/contributing_link>
 Contributors <misc/contributors_link>
 Known Issues <known_issues>
 Changelog <misc/changelog_link>
+References <misc/references>
 Github <https://github.com/emdgroup/baybe/>
 License <misc/license_link>
 ```
diff --git a/docs/misc/references.md b/docs/misc/references.md
new file mode 100644
index 0000000000..d977a9bd5b
--- /dev/null
+++ b/docs/misc/references.md
@@ -0,0 +1,4 @@
+# References
+
+```{bibliography}
+```
diff --git a/docs/references.bib b/docs/references.bib
index 307529d299..4d159cb686 100644
--- a/docs/references.bib
+++ b/docs/references.bib
@@ -8,3 +8,22 @@ @inproceedings{NIPS2007_66368270
  volume = {20},
  year = {2007}
 }
+
+@article{Chen2026,
+ author = {Chen, Guanming and Fleck, Maximilian and Stuyver, Thijs},
+ title = {Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors},
+ journal = {ChemRxiv},
+ year = {2026},
+ doi = {10.26434/chemrxiv.10001986/v2}
+}
+
+@article{Shields2021,
+ author = {Shields, Benjamin J. and Stevens, Jason and Li, Jun and Parasram, Marvin and Damani, Farhan and Alvarado, Jesus I. Martinez and Janey, Jacob M. and Adams, Ryan P. and Doyle, Abigail G.},
+ title = {Bayesian reaction optimization as a tool for chemical synthesis},
+ journal = {Nature},
+ volume = {590},
+ number = {7844},
+ pages = {89--96},
+ year = {2021},
+ doi = {10.1038/s41586-021-03213-y}
+}
diff --git a/docs/userguide/transfer_learning.md b/docs/userguide/transfer_learning.md
index 6f509df31b..3273e20853 100644
--- a/docs/userguide/transfer_learning.md
+++ b/docs/userguide/transfer_learning.md
@@ -185,7 +185,4 @@ on the optimization:
 :class: only-dark
 ```
 
-```{bibliography}
-```
-
 [`TaskParameter`]: baybe.parameters.categorical.TaskParameter
\ No newline at end of file

From 6a50f5df202807f38812261e6ee6093fb272fbb2 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 16 Apr 2026 13:28:24 +0200
Subject: [PATCH 15/16] Update GaussianProcessPreset enum

---
 baybe/surrogates/gaussian_process/presets/core.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py
index ad77df0b4d..5347cf85e5 100644
--- a/baybe/surrogates/gaussian_process/presets/core.py
+++ b/baybe/surrogates/gaussian_process/presets/core.py
@@ -11,8 +11,11 @@ class GaussianProcessPreset(Enum):
     BAYBE = "BAYBE"
     """The default BayBE settings of the Gaussian process surrogate class."""
 
+    CHEN = "CHEN"
+    """The adaptive kernel hyperprior settings proposed by :cite:p:`Chen2026`."""
+
     EDBO = "EDBO"
-    """The EDBO settings."""
+    """The EDBO settings proposed by :cite:p:`Shields2021`."""
 
     EDBO_SMOOTHED = "EDBO_SMOOTHED"
-    """A smoothed version of the EDBO settings."""
+    """A smoothed version of the EDBO settings (adapted from :cite:p:`Shields2021`)."""

From 2398672dbc204725b9b6325cd0f5f94412a3a713 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 16 Apr 2026 13:29:28 +0200
Subject: [PATCH 16/16] Update CHANGELOG.md

---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c7715778a3..de3165f616 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support for GPyTorch objects (kernels, means, likelihood) as Gaussian process
   components, enabling full low-level customization
 - Factories for all Gaussian process components
-- `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate`
+- `CHEN`, `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate`
 - `TypeSelector` and `NameSelector` classes for parameter selection in kernel factories
 - `parameter_names` attribute to basic kernels for controlling the considered parameters
 - `ParameterKind` flag enum for classifying parameters by their role and automatic