From 0611ff6e944d5d72ce6e54da99aafdff11b1d706 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 09:48:53 +0200
Subject: [PATCH 01/13] Extract transfer learning mechanism into a reusable
 decorator

* Provides a single source of truth for defining the TL logic
* Enables TL for non-TL presets by applying the decorator
---
 .../gaussian_process/components/kernel.py     | 57 +++++++++++++++++--
 .../gaussian_process/presets/baybe.py         | 27 ++-------
 .../gaussian_process/presets/chen.py          | 16 +-----
 .../gaussian_process/presets/edbo.py          | 15 +----
 .../gaussian_process/presets/edbo_smoothed.py | 33 +++++------
 tests/test_kernel_factories.py                | 12 ++--
 6 files changed, 85 insertions(+), 75 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index 5bc0415da7..c39bcd005f 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
 
+import functools
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from functools import partial
@@ -115,6 +116,47 @@ def _make(
         """Construct the kernel."""
 
 
+def _enable_transfer_learning(
+    cls: type[_PureKernelFactory], /
+) -> type[_PureKernelFactory]:
+    """Class decorator enabling BayBE's default transfer learning mechanism.
+
+    When the search space contains a task parameter, the decorated factory
+    automatically composes its kernel with BayBE's default task kernel.
+    Otherwise, the factory behaves unchanged.
+
+    Args:
+        cls: The kernel factory class to decorate.
+
+    Raises:
+        TypeError: If the factory already supports task parameters.
+
+    Returns:
+        The decorated kernel factory class with transfer learning enabled.
+    """
+    if cls._supported_parameter_kinds & _ParameterKind.TASK:
+        raise TypeError(f"'{cls.__name__}' already supports task parameters.")
+
+    # Create a subclass so the original class remains unmodified
+    new_cls = type(cls.__name__, (cls,), {"__doc__": cls.__doc__})
+
+    original_call = cls.__call__
+
+    @functools.wraps(original_call)
+    def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
+        base_kernel = original_call(self, searchspace, train_x, train_y)
+        if searchspace.task_idx is not None:
+            icm = ICMKernelFactory(base_kernel_or_factory=base_kernel)
+            return icm(searchspace, train_x, train_y)
+        return base_kernel
+
+    new_cls.__call__ = __call__  # type: ignore[method-assign]
+    new_cls._supported_parameter_kinds = (
+        cls._supported_parameter_kinds | _ParameterKind.TASK
+    )
+    return new_cls
+
+
 @define
 class _MetaKernelFactory(KernelFactoryProtocol, ABC):
     """Base class for meta kernel factories that orchestrate other kernel factories."""
@@ -150,18 +192,25 @@ class ICMKernelFactory(_MetaKernelFactory):
     @base_kernel_factory.default
     def _default_base_kernel_factory(self) -> KernelFactoryProtocol:
         from baybe.surrogates.gaussian_process.presets.baybe import (
-            BayBENumericalKernelFactory,
+            _BayBENumericalKernelFactory,
         )
 
-        return BayBENumericalKernelFactory(TypeSelector((TaskParameter,), exclude=True))
+        assert (
+            _BayBENumericalKernelFactory._supported_parameter_kinds
+            is _ParameterKind.REGULAR
+        )
+        return _BayBENumericalKernelFactory(
+            TypeSelector((TaskParameter,), exclude=True)
+        )
 
     @task_kernel_factory.default
     def _default_task_kernel_factory(self) -> KernelFactoryProtocol:
         from baybe.surrogates.gaussian_process.presets.baybe import (
-            BayBETaskKernelFactory,
+            _BayBETaskKernelFactory,
         )
 
-        return BayBETaskKernelFactory(TypeSelector((TaskParameter,)))
+        assert _BayBETaskKernelFactory._supported_parameter_kinds is _ParameterKind.TASK
+        return _BayBETaskKernelFactory()
 
     @override
     def __call__(
diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py
index 0cf6686064..97eaaa8021 100644
--- a/baybe/surrogates/gaussian_process/presets/baybe.py
+++ b/baybe/surrogates/gaussian_process/presets/baybe.py
@@ -26,38 +26,23 @@
 from baybe.surrogates.gaussian_process.presets.edbo_smoothed import (
     SmoothedEDBOKernelFactory,
     SmoothedEDBOLikelihoodFactory,
+    _SmoothedEDBONumericalKernelFactory,
 )
 
 if TYPE_CHECKING:
     from torch import Tensor
 
 
-@define
-class BayBEKernelFactory(_PureKernelFactory):
-    """The default kernel factory for Gaussian process surrogates."""
-
-    _supported_parameter_kinds: ClassVar[_ParameterKind] = (
-        _ParameterKind.REGULAR | _ParameterKind.TASK
-    )
-    # See base class.
-
-    @override
-    def _make(
-        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel:
-        from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory
-
-        is_multitask = searchspace.task_idx is not None
-        factory = ICMKernelFactory if is_multitask else BayBENumericalKernelFactory
-        return factory()(searchspace, train_x, train_y)
+_BayBENumericalKernelFactory = _SmoothedEDBONumericalKernelFactory
+"""The factory providing the default numerical kernel for Gaussian process surrogates."""  # noqa: E501
 
 
-BayBENumericalKernelFactory = SmoothedEDBOKernelFactory
-"""The factory providing the default numerical kernel for Gaussian process surrogates."""  # noqa: E501
+BayBEKernelFactory = SmoothedEDBOKernelFactory
+"""The default kernel factory for Gaussian process surrogates."""
 
 
 @define
-class BayBETaskKernelFactory(_PureKernelFactory):
+class _BayBETaskKernelFactory(_PureKernelFactory):
     """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/chen.py b/baybe/surrogates/gaussian_process/presets/chen.py
index e461bc0750..cf234246d3 100644
--- a/baybe/surrogates/gaussian_process/presets/chen.py
+++ b/baybe/surrogates/gaussian_process/presets/chen.py
@@ -6,22 +6,17 @@
 import math
 from typing import TYPE_CHECKING, ClassVar
 
-from attrs import define, field
+from attrs import define
 from typing_extensions import override
 
 from baybe.kernels.basic import MaternKernel
 from baybe.kernels.composite import ScaleKernel
-from baybe.parameters.categorical import TaskParameter
-from baybe.parameters.selectors import (
-    ParameterSelectorProtocol,
-    TypeSelector,
-    to_parameter_selector,
-)
 from baybe.priors.basic import GammaPrior
 from baybe.surrogates.gaussian_process.components.fit_criterion import (
     _MLLForNonTLFitCriterionFactory,
 )
 from baybe.surrogates.gaussian_process.components.kernel import (
+    _enable_transfer_learning,
     _PureKernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
@@ -36,6 +31,7 @@
     from baybe.searchspace.core import SearchSpace
 
 
+@_enable_transfer_learning
 @define
 class CHENKernelFactory(_PureKernelFactory):
     """A factory providing adaptive hyperprior kernels as proposed by :cite:p:`Chen2026`."""  # noqa: E501
@@ -43,12 +39,6 @@ class CHENKernelFactory(_PureKernelFactory):
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
 
-    parameter_selector: ParameterSelectorProtocol | None = field(
-        factory=lambda: TypeSelector([TaskParameter], exclude=True),
-        converter=to_parameter_selector,
-    )
-    # TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
-
     @override
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 1ff8d2bf80..7157af9c2e 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -6,18 +6,13 @@
 from collections.abc import Collection
 from typing import TYPE_CHECKING, ClassVar
 
-from attrs import define, field
+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.parameters.enum import SubstanceEncoding
-from baybe.parameters.selectors import (
-    ParameterSelectorProtocol,
-    TypeSelector,
-    to_parameter_selector,
-)
 from baybe.parameters.substance import SubstanceParameter
 from baybe.priors.basic import GammaPrior
 from baybe.searchspace.discrete import SubspaceDiscrete
@@ -25,6 +20,7 @@
     _MLLForNonTLFitCriterionFactory,
 )
 from baybe.surrogates.gaussian_process.components.kernel import (
+    _enable_transfer_learning,
     _PureKernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
@@ -59,6 +55,7 @@ def _contains_encoding(
 """Encodings relevant to EDBO logic."""
 
 
+@_enable_transfer_learning
 @define
 class EDBOKernelFactory(_PureKernelFactory):
     """A factory providing EDBO kernels, as proposed by :cite:p:`Shields2021`.
@@ -70,12 +67,6 @@ class EDBOKernelFactory(_PureKernelFactory):
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
 
-    parameter_selector: ParameterSelectorProtocol | None = field(
-        factory=lambda: TypeSelector([TaskParameter], exclude=True),
-        converter=to_parameter_selector,
-    )
-    # TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
-
     @override
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index 519713f9dc..1684a0c4d1 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -6,22 +6,18 @@
 from typing import TYPE_CHECKING, ClassVar
 
 import numpy as np
-from attrs import define, field
+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.parameters.selectors import (
-    ParameterSelectorProtocol,
-    TypeSelector,
-    to_parameter_selector,
-)
 from baybe.priors.basic import GammaPrior
 from baybe.surrogates.gaussian_process.components.fit_criterion import (
     _MLLForNonTLFitCriterionFactory,
 )
 from baybe.surrogates.gaussian_process.components.kernel import (
+    _enable_transfer_learning,
     _PureKernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
@@ -42,23 +38,12 @@
 
 
 @define
-class SmoothedEDBOKernelFactory(_PureKernelFactory):
-    """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
+class _SmoothedEDBONumericalKernelFactory(_PureKernelFactory):
+    """A factory providing the core numerical kernel for the smoothed EDBO preset."""
 
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
 
-    parameter_selector: ParameterSelectorProtocol | None = field(
-        factory=lambda: TypeSelector([TaskParameter], exclude=True),
-        converter=to_parameter_selector,
-    )
-    # TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
-
     @override
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
@@ -91,6 +76,16 @@ def _make(
         )
 
 
+SmoothedEDBOKernelFactory = _enable_transfer_learning(
+    _SmoothedEDBONumericalKernelFactory
+)
+"""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
+
 SmoothedEDBOMeanFactory = LazyConstantMeanFactory
 """A factory providing mean functions for the smoothed EDBO preset."""
 
diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py
index 6a8acda6a6..45f3ccc109 100644
--- a/tests/test_kernel_factories.py
+++ b/tests/test_kernel_factories.py
@@ -15,8 +15,8 @@
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.gaussian_process.presets.baybe import (
     BayBEKernelFactory,
-    BayBENumericalKernelFactory,
-    BayBETaskKernelFactory,
+    _BayBENumericalKernelFactory,
+    _BayBETaskKernelFactory,
 )
 
 # A selector that accepts all parameters
@@ -27,25 +27,25 @@
     ("factory", "parameters", "error"),
     [
         param(
-            BayBENumericalKernelFactory(parameter_selector=_SELECT_ALL),
+            _BayBENumericalKernelFactory(parameter_selector=_SELECT_ALL),
             [TaskParameter("task", ["t1", "t2"])],
             IncompatibleSearchSpaceError,
             id="regular_rejects_task",
         ),
         param(
-            BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
+            _BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
             [CategoricalParameter("cat", ["a", "b"])],
             IncompatibleSearchSpaceError,
             id="task_rejects_categorical",
         ),
         param(
-            BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
+            _BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
             [NumericalDiscreteParameter("num", [1, 2, 3])],
             IncompatibleSearchSpaceError,
             id="task_rejects_numerical_discrete",
         ),
         param(
-            BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
+            _BayBETaskKernelFactory(parameter_selector=_SELECT_ALL),
             [NumericalContinuousParameter("cont", (0, 1))],
             IncompatibleSearchSpaceError,
             id="task_rejects_numerical_continuous",

From 43d1e8e599417589764d9142304e81a0cb11b2ff Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 16:01:01 +0200
Subject: [PATCH 02/13] Scope inner factory to non-task parameters in transfer
 learning decorator

---
 .../gaussian_process/components/kernel.py     | 32 ++++++++++++++++++-
 .../gaussian_process/presets/chen.py          |  3 +-
 .../gaussian_process/presets/edbo.py          | 11 ++++---
 .../gaussian_process/presets/edbo_smoothed.py | 10 +++---
 4 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index c39bcd005f..febcd9f1dc 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -76,6 +76,15 @@ def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...]:
         selector = self.parameter_selector or (lambda _: True)
         return tuple(p.name for p in searchspace.parameters if selector(p))
 
+    def _get_effective_dimensionality(self, searchspace: SearchSpace) -> int:
+        """Get the number of computational columns for the selected parameters."""
+        names = self.get_parameter_names(searchspace)
+        if names is None:
+            return len(searchspace.comp_rep_columns)
+        return sum(
+            len(searchspace.get_comp_rep_parameter_indices(name)) for name in names
+        )
+
     def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None:
         """Validate that the given parameters are supported by the factory.
 
@@ -141,10 +150,31 @@ def _enable_transfer_learning(
     new_cls = type(cls.__name__, (cls,), {"__doc__": cls.__doc__})
 
     original_call = cls.__call__
+    original_supported_kinds = cls._supported_parameter_kinds
+    _task_exclude_selector = TypeSelector((TaskParameter,), exclude=True)
 
     @functools.wraps(original_call)
     def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
-        base_kernel = original_call(self, searchspace, train_x, train_y)
+        # Temporarily narrow the supported parameter kinds to those of the original
+        # class. If the decorator logic is correct, the original factory should never
+        # see the extended scope, but this acts as a sanity check to prevent regressions
+        broadened_kinds = new_cls._supported_parameter_kinds
+        new_cls._supported_parameter_kinds = original_supported_kinds
+
+        # Split off the task parameters
+        original_selector = self.parameter_selector
+        if original_selector is None:
+            self.parameter_selector = _task_exclude_selector
+        else:
+            self.parameter_selector = lambda p: (
+                _task_exclude_selector(p) and original_selector(p)
+            )
+        try:
+            base_kernel = original_call(self, searchspace, train_x, train_y)
+        finally:
+            new_cls._supported_parameter_kinds = broadened_kinds
+            self.parameter_selector = original_selector
+
         if searchspace.task_idx is not None:
             icm = ICMKernelFactory(base_kernel_or_factory=base_kernel)
             return icm(searchspace, train_x, train_y)
diff --git a/baybe/surrogates/gaussian_process/presets/chen.py b/baybe/surrogates/gaussian_process/presets/chen.py
index cf234246d3..0122004beb 100644
--- a/baybe/surrogates/gaussian_process/presets/chen.py
+++ b/baybe/surrogates/gaussian_process/presets/chen.py
@@ -43,7 +43,8 @@ class CHENKernelFactory(_PureKernelFactory):
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
-        lengthscale = 0.4 * math.sqrt(train_x.shape[-1]) + 4.0
+        n_dimensions = self._get_effective_dimensionality(searchspace)
+        lengthscale = 0.4 * math.sqrt(n_dimensions) + 4.0
         lengthscale_prior = GammaPrior(2.0 * lengthscale, 2.0)
         lengthscale_initial_value = lengthscale
         outputscale_prior = GammaPrior(1.0 * lengthscale, 1.0)
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 7157af9c2e..3657ce4e65 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -11,8 +11,7 @@
 
 from baybe.kernels.basic import MaternKernel
 from baybe.kernels.composite import ScaleKernel
-from baybe.parameters import TaskParameter
-from baybe.parameters.enum import SubstanceEncoding
+from baybe.parameters.enum import SubstanceEncoding, _ParameterKind
 from baybe.parameters.substance import SubstanceParameter
 from baybe.priors.basic import GammaPrior
 from baybe.searchspace.discrete import SubspaceDiscrete
@@ -71,7 +70,7 @@ class EDBOKernelFactory(_PureKernelFactory):
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
-        effective_dims = train_x.shape[-1]
+        effective_dims = self._get_effective_dimensionality(searchspace)
 
         switching_condition = _contains_encoding(
             searchspace.discrete, _EDBO_ENCODINGS
@@ -136,8 +135,10 @@ def __call__(
         import torch
         from gpytorch.likelihoods import GaussianLikelihood
 
-        effective_dims = train_x.shape[-1] - len(
-            [p for p in searchspace.parameters if isinstance(p, TaskParameter)]
+        effective_dims = sum(
+            len(searchspace.get_comp_rep_parameter_indices(p.name))
+            for p in searchspace.parameters
+            if p._kind & _ParameterKind.REGULAR
         )
 
         switching_condition = _contains_encoding(
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index 1684a0c4d1..82a6927d3b 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -11,7 +11,7 @@
 
 from baybe.kernels.basic import MaternKernel
 from baybe.kernels.composite import ScaleKernel
-from baybe.parameters import TaskParameter
+from baybe.parameters.enum import _ParameterKind
 from baybe.priors.basic import GammaPrior
 from baybe.surrogates.gaussian_process.components.fit_criterion import (
     _MLLForNonTLFitCriterionFactory,
@@ -48,7 +48,7 @@ class _SmoothedEDBONumericalKernelFactory(_PureKernelFactory):
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
-        effective_dims = train_x.shape[-1]
+        effective_dims = self._get_effective_dimensionality(searchspace)
 
         # 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.
@@ -109,8 +109,10 @@ def __call__(
         # 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.
         # Values outside the dimension limits will get the border value assigned.
-        effective_dims = train_x.shape[-1] - len(
-            [p for p in searchspace.parameters if isinstance(p, TaskParameter)]
+        effective_dims = sum(
+            len(searchspace.get_comp_rep_parameter_indices(p.name))
+            for p in searchspace.parameters
+            if p._kind & _ParameterKind.REGULAR
         )
 
         prior = GammaPrior(

From 542029bd530fd4d8cf9483378bced5b3d91c57ae Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 16:44:34 +0200
Subject: [PATCH 03/13] Fix class name of dynamically created kernel factory

`_enable_transfer_learning` now accepts an optional `name` parameter
so that the dynamically created class can have the correct `__name__`
when the function is called directly (rather than used as a decorator).
This fixes serialization for `SmoothedEDBOKernelFactory`, which was
previously serialized as `_SmoothedEDBONumericalKernelFactory`.
---
 baybe/surrogates/gaussian_process/components/kernel.py     | 7 +++++--
 baybe/surrogates/gaussian_process/presets/edbo_smoothed.py | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index febcd9f1dc..067b2a898c 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -126,7 +126,7 @@ def _make(
 
 
 def _enable_transfer_learning(
-    cls: type[_PureKernelFactory], /
+    cls: type[_PureKernelFactory], name: str | None = None, /
 ) -> type[_PureKernelFactory]:
     """Class decorator enabling BayBE's default transfer learning mechanism.
 
@@ -136,6 +136,9 @@ def _enable_transfer_learning(
 
     Args:
         cls: The kernel factory class to decorate.
+        name: Optional name for the created class. Defaults to ``cls.__name__``.
+            Useful when calling the function directly (as opposed to using it as a
+            decorator) and assigning the result to a different name.
 
     Raises:
         TypeError: If the factory already supports task parameters.
@@ -147,7 +150,7 @@ def _enable_transfer_learning(
         raise TypeError(f"'{cls.__name__}' already supports task parameters.")
 
     # Create a subclass so the original class remains unmodified
-    new_cls = type(cls.__name__, (cls,), {"__doc__": cls.__doc__})
+    new_cls = type(name or cls.__name__, (cls,), {"__doc__": cls.__doc__})
 
     original_call = cls.__call__
     original_supported_kinds = cls._supported_parameter_kinds
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index 82a6927d3b..4564986508 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -77,7 +77,7 @@ def _make(
 
 
 SmoothedEDBOKernelFactory = _enable_transfer_learning(
-    _SmoothedEDBONumericalKernelFactory
+    _SmoothedEDBONumericalKernelFactory, "SmoothedEDBOKernelFactory"
 )
 """A factory providing smoothed versions of EDBO kernels (adapted from :cite:p:`Shields2021`).
 

From e41a78d99ebc5da99617c70eefad81e3940e3650 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 17:41:55 +0200
Subject: [PATCH 04/13] Replace factory aliases with thin subclasses

Simple aliases like `BayBEKernelFactory = SmoothedEDBOKernelFactory`
cause the serialized type name to be that of the underlying class,
which means the identity is lost on deserialization. Using thin
subclasses ensures each factory has its own stable `__name__`.
---
 .../gaussian_process/presets/baybe.py         | 20 ++++++++++---------
 .../gaussian_process/presets/edbo.py          |  4 ++--
 .../gaussian_process/presets/edbo_smoothed.py |  5 +++--
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py
index 97eaaa8021..f59e733968 100644
--- a/baybe/surrogates/gaussian_process/presets/baybe.py
+++ b/baybe/surrogates/gaussian_process/presets/baybe.py
@@ -33,17 +33,17 @@
     from torch import Tensor
 
 
-_BayBENumericalKernelFactory = _SmoothedEDBONumericalKernelFactory
-"""The factory providing the default numerical kernel for Gaussian process surrogates."""  # noqa: E501
+class _BayBENumericalKernelFactory(_SmoothedEDBONumericalKernelFactory):
+    """The default numerical kernel factory for GP surrogates."""
 
 
-BayBEKernelFactory = SmoothedEDBOKernelFactory
-"""The default kernel factory for Gaussian process surrogates."""
+class BayBEKernelFactory(SmoothedEDBOKernelFactory):
+    """The default kernel factory for GP surrogates."""
 
 
 @define
 class _BayBETaskKernelFactory(_PureKernelFactory):
-    """The factory providing the default task kernel for Gaussian process surrogates."""
+    """The default task kernel factory for GP surrogates."""
 
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
@@ -68,11 +68,13 @@ def _make(
         )
 
 
-BayBEMeanFactory = LazyConstantMeanFactory
-"""The factory providing the default mean function for Gaussian process surrogates."""
+class BayBEMeanFactory(LazyConstantMeanFactory):
+    """The default mean factory for GP surrogates."""
+
+
+class BayBELikelihoodFactory(SmoothedEDBOLikelihoodFactory):
+    """The default likelihood factory for GP surrogates."""
 
-BayBELikelihoodFactory = SmoothedEDBOLikelihoodFactory
-"""The factory providing the default likelihood for Gaussian process surrogates."""
 
 
 @define
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 3657ce4e65..2e8596404f 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -116,8 +116,8 @@ def _make(
         )
 
 
-EDBOMeanFactory = LazyConstantMeanFactory
-"""A factory providing mean functions for the EDBO preset."""
+class EDBOMeanFactory(LazyConstantMeanFactory):
+    """A factory providing mean functions for the EDBO preset."""
 
 
 @define
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index 4564986508..2a9dbb441f 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -86,8 +86,9 @@ def _make(
 and interpolates the prior moments linearly in between.
 """  # noqa: E501
 
-SmoothedEDBOMeanFactory = LazyConstantMeanFactory
-"""A factory providing mean functions for the smoothed EDBO preset."""
+
+class SmoothedEDBOMeanFactory(LazyConstantMeanFactory):
+    """A factory providing mean functions for the smoothed EDBO preset."""
 
 
 @define

From ded217f335dc982888d65e1849bee59e93013eeb Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 18:42:48 +0200
Subject: [PATCH 05/13] Add serialization roundtrip tests for kernel factories

---
 tests/test_kernel_factories.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py
index 45f3ccc109..6c87e63c9f 100644
--- a/tests/test_kernel_factories.py
+++ b/tests/test_kernel_factories.py
@@ -13,11 +13,17 @@
     NumericalDiscreteParameter,
 )
 from baybe.searchspace.core import SearchSpace
+from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate
 from baybe.surrogates.gaussian_process.presets.baybe import (
     BayBEKernelFactory,
     _BayBENumericalKernelFactory,
     _BayBETaskKernelFactory,
 )
+from baybe.surrogates.gaussian_process.presets.chen import CHENKernelFactory
+from baybe.surrogates.gaussian_process.presets.edbo import EDBOKernelFactory
+from baybe.surrogates.gaussian_process.presets.edbo_smoothed import (
+    SmoothedEDBOKernelFactory,
+)
 
 # A selector that accepts all parameters
 _SELECT_ALL = lambda parameter: True  # noqa: E731
@@ -73,3 +79,21 @@ def test_factory_parameter_kind_validation(factory, parameters, error):
         else pytest.raises(error, match="does not support")
     ):
         factory(ss, train_x, train_y)
+
+
+@pytest.mark.parametrize(
+    "factory",
+    [
+        param(BayBEKernelFactory(), id="BayBEKernelFactory"),
+        param(SmoothedEDBOKernelFactory(), id="SmoothedEDBOKernelFactory"),
+        param(EDBOKernelFactory(), id="EDBOKernelFactory"),
+        param(CHENKernelFactory(), id="CHENKernelFactory"),
+    ],
+)
+def test_kernel_factory_serialization_roundtrip(factory):
+    """Kernel factories survive a serialization roundtrip via a GP surrogate."""
+    gp = GaussianProcessSurrogate(kernel_or_factory=factory)
+    json_str = gp.to_json()
+    gp_roundtrip = GaussianProcessSurrogate.from_json(json_str)
+    assert type(gp.kernel_factory) is type(gp_roundtrip.kernel_factory)
+    assert gp == gp_roundtrip

From 05b89b262cf378c3ecc005455895b0917718fb82 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 20:09:43 +0200
Subject: [PATCH 06/13] Fix serialization of transfer-learning-decorated kernel
 factories

When used as a decorator (@_enable_transfer_learning), modify the class
in-place instead of creating a subclass with the same __name__. The
previous approach left two concrete classes with identical names in the
subclass registry, causing find_subclass to resolve to the @define-
processed intermediate (without the TL wrapper) during deserialization.

When called with an explicit name argument (for cases like
SmoothedEDBOKernelFactory where the original class is reused elsewhere),
the subclass approach is preserved since the distinct name avoids any
collision.
---
 .../gaussian_process/components/kernel.py     | 32 ++++++++++++-------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index 067b2a898c..be441187e6 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -134,11 +134,15 @@ def _enable_transfer_learning(
     automatically composes its kernel with BayBE's default task kernel.
     Otherwise, the factory behaves unchanged.
 
+    When used as a decorator (without ``name``), the class is modified in-place.
+    When called with a ``name`` argument, a new subclass is created so that the
+    original class remains unmodified. The latter form is intended for cases where
+    the original class is reused independently elsewhere.
+
     Args:
         cls: The kernel factory class to decorate.
-        name: Optional name for the created class. Defaults to ``cls.__name__``.
-            Useful when calling the function directly (as opposed to using it as a
-            decorator) and assigning the result to a different name.
+        name: Optional name for the created class. If provided, a new subclass is
+            created instead of modifying ``cls`` in-place.
 
     Raises:
         TypeError: If the factory already supports task parameters.
@@ -149,8 +153,14 @@ def _enable_transfer_learning(
     if cls._supported_parameter_kinds & _ParameterKind.TASK:
         raise TypeError(f"'{cls.__name__}' already supports task parameters.")
 
-    # Create a subclass so the original class remains unmodified
-    new_cls = type(name or cls.__name__, (cls,), {"__doc__": cls.__doc__})
+    # This distinction is important for serialization so that the classes can be
+    # correctly identified by their names in the subclass registry
+    if name is not None:
+        # Create a subclass so the original class remains unmodified
+        target_cls = type(name, (cls,), {"__doc__": cls.__doc__})
+    else:
+        # Modify the class in-place (avoids name collision in subclass registry)
+        target_cls = cls
 
     original_call = cls.__call__
     original_supported_kinds = cls._supported_parameter_kinds
@@ -161,8 +171,8 @@ def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
         # Temporarily narrow the supported parameter kinds to those of the original
         # class. If the decorator logic is correct, the original factory should never
         # see the extended scope, but this acts as a sanity check to prevent regressions
-        broadened_kinds = new_cls._supported_parameter_kinds
-        new_cls._supported_parameter_kinds = original_supported_kinds
+        broadened_kinds = target_cls._supported_parameter_kinds
+        target_cls._supported_parameter_kinds = original_supported_kinds
 
         # Split off the task parameters
         original_selector = self.parameter_selector
@@ -175,7 +185,7 @@ def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
         try:
             base_kernel = original_call(self, searchspace, train_x, train_y)
         finally:
-            new_cls._supported_parameter_kinds = broadened_kinds
+            target_cls._supported_parameter_kinds = broadened_kinds
             self.parameter_selector = original_selector
 
         if searchspace.task_idx is not None:
@@ -183,11 +193,11 @@ def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
             return icm(searchspace, train_x, train_y)
         return base_kernel
 
-    new_cls.__call__ = __call__  # type: ignore[method-assign]
-    new_cls._supported_parameter_kinds = (
+    target_cls.__call__ = __call__  # type: ignore[method-assign]
+    target_cls._supported_parameter_kinds = (
         cls._supported_parameter_kinds | _ParameterKind.TASK
     )
-    return new_cls
+    return target_cls
 
 
 @define

From 21ed53719e63be59e9bb66ae00080c6d54cbe540 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 20:16:05 +0200
Subject: [PATCH 07/13] Add SerialMixin to _PureKernelFactory and move
 serialization test

---
 .../gaussian_process/components/kernel.py     |  3 ++-
 .../test_kernel_factory_serialization.py      | 23 ++++++++++++++++++
 tests/test_kernel_factories.py                | 24 -------------------
 3 files changed, 25 insertions(+), 25 deletions(-)
 create mode 100644 tests/serialization/test_kernel_factory_serialization.py

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index be441187e6..faee77ef33 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -23,6 +23,7 @@
     to_parameter_selector,
 )
 from baybe.searchspace.core import SearchSpace
+from baybe.serialization.mixin import SerialMixin
 from baybe.surrogates.gaussian_process.components.generic import (
     GPComponentFactoryProtocol,
     GPComponentType,
@@ -45,7 +46,7 @@
 
 
 @define
-class _PureKernelFactory(KernelFactoryProtocol, ABC):
+class _PureKernelFactory(KernelFactoryProtocol, SerialMixin, ABC):
     """Base class for pure kernel factories."""
 
     # For internal use only: sanity check mechanism to remind developers of new
diff --git a/tests/serialization/test_kernel_factory_serialization.py b/tests/serialization/test_kernel_factory_serialization.py
new file mode 100644
index 0000000000..0bdc20af93
--- /dev/null
+++ b/tests/serialization/test_kernel_factory_serialization.py
@@ -0,0 +1,23 @@
+"""Kernel factory serialization tests."""
+
+import pytest
+
+from baybe.surrogates.gaussian_process.components.kernel import _PureKernelFactory
+from baybe.surrogates.gaussian_process.presets import *  # noqa: F401, F403
+from baybe.utils.basic import get_subclasses
+from tests.serialization.utils import assert_roundtrip_consistency
+
+_KERNEL_FACTORIES = [
+    cls
+    for cls in get_subclasses(_PureKernelFactory)
+    if not cls.__name__.startswith("_")
+]
+
+
+@pytest.mark.parametrize(
+    "factory",
+    [pytest.param(cls(), id=cls.__name__) for cls in _KERNEL_FACTORIES],
+)
+def test_roundtrip(factory):
+    """A serialization roundtrip yields an equivalent object."""
+    assert_roundtrip_consistency(factory)
diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py
index 6c87e63c9f..45f3ccc109 100644
--- a/tests/test_kernel_factories.py
+++ b/tests/test_kernel_factories.py
@@ -13,17 +13,11 @@
     NumericalDiscreteParameter,
 )
 from baybe.searchspace.core import SearchSpace
-from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate
 from baybe.surrogates.gaussian_process.presets.baybe import (
     BayBEKernelFactory,
     _BayBENumericalKernelFactory,
     _BayBETaskKernelFactory,
 )
-from baybe.surrogates.gaussian_process.presets.chen import CHENKernelFactory
-from baybe.surrogates.gaussian_process.presets.edbo import EDBOKernelFactory
-from baybe.surrogates.gaussian_process.presets.edbo_smoothed import (
-    SmoothedEDBOKernelFactory,
-)
 
 # A selector that accepts all parameters
 _SELECT_ALL = lambda parameter: True  # noqa: E731
@@ -79,21 +73,3 @@ def test_factory_parameter_kind_validation(factory, parameters, error):
         else pytest.raises(error, match="does not support")
     ):
         factory(ss, train_x, train_y)
-
-
-@pytest.mark.parametrize(
-    "factory",
-    [
-        param(BayBEKernelFactory(), id="BayBEKernelFactory"),
-        param(SmoothedEDBOKernelFactory(), id="SmoothedEDBOKernelFactory"),
-        param(EDBOKernelFactory(), id="EDBOKernelFactory"),
-        param(CHENKernelFactory(), id="CHENKernelFactory"),
-    ],
-)
-def test_kernel_factory_serialization_roundtrip(factory):
-    """Kernel factories survive a serialization roundtrip via a GP surrogate."""
-    gp = GaussianProcessSurrogate(kernel_or_factory=factory)
-    json_str = gp.to_json()
-    gp_roundtrip = GaussianProcessSurrogate.from_json(json_str)
-    assert type(gp.kernel_factory) is type(gp_roundtrip.kernel_factory)
-    assert gp == gp_roundtrip

From 86ff984034ba8717d136d1e4344841a510556e0e Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 20:29:35 +0200
Subject: [PATCH 08/13] Fix __module__ on dynamically-created kernel factory
 subclass

The Protocol metaclass (_ProtocolMeta) defaults __module__ to 'abc'
when creating classes via 3-arg type(). Set it explicitly from the
parent class so that SmoothedEDBOKernelFactory correctly reports its
module as baybe.surrogates.gaussian_process.presets.edbo_smoothed.
---
 baybe/surrogates/gaussian_process/components/kernel.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index faee77ef33..debf60e36f 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -157,8 +157,12 @@ def _enable_transfer_learning(
     # This distinction is important for serialization so that the classes can be
     # correctly identified by their names in the subclass registry
     if name is not None:
-        # Create a subclass so the original class remains unmodified
-        target_cls = type(name, (cls,), {"__doc__": cls.__doc__})
+        # Create a subclass so the original class remains unmodified.
+        # __module__ must be set explicitly because the Protocol metaclass
+        # would otherwise default it to "abc".
+        target_cls = type(
+            name, (cls,), {"__doc__": cls.__doc__, "__module__": cls.__module__}
+        )
     else:
         # Modify the class in-place (avoids name collision in subclass registry)
         target_cls = cls

From 41a1fc7334569f302f71e63d93066fc2ddb4a1ca Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Thu, 7 May 2026 20:35:30 +0200
Subject: [PATCH 09/13] Suppress mypy errors from dynamic class creation in
 _enable_transfer_learning

---
 baybe/surrogates/gaussian_process/components/kernel.py | 8 ++++----
 baybe/surrogates/gaussian_process/presets/baybe.py     | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index debf60e36f..ecb43f0a10 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -176,8 +176,8 @@ def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
         # Temporarily narrow the supported parameter kinds to those of the original
         # class. If the decorator logic is correct, the original factory should never
         # see the extended scope, but this acts as a sanity check to prevent regressions
-        broadened_kinds = target_cls._supported_parameter_kinds
-        target_cls._supported_parameter_kinds = original_supported_kinds
+        broadened_kinds = target_cls._supported_parameter_kinds  # type: ignore[attr-defined]
+        target_cls._supported_parameter_kinds = original_supported_kinds  # type: ignore[attr-defined]
 
         # Split off the task parameters
         original_selector = self.parameter_selector
@@ -190,7 +190,7 @@ def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
         try:
             base_kernel = original_call(self, searchspace, train_x, train_y)
         finally:
-            target_cls._supported_parameter_kinds = broadened_kinds
+            target_cls._supported_parameter_kinds = broadened_kinds  # type: ignore[attr-defined]
             self.parameter_selector = original_selector
 
         if searchspace.task_idx is not None:
@@ -199,7 +199,7 @@ def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
         return base_kernel
 
     target_cls.__call__ = __call__  # type: ignore[method-assign]
-    target_cls._supported_parameter_kinds = (
+    target_cls._supported_parameter_kinds = (  # type: ignore[attr-defined]
         cls._supported_parameter_kinds | _ParameterKind.TASK
     )
     return target_cls
diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py
index f59e733968..520fb00acd 100644
--- a/baybe/surrogates/gaussian_process/presets/baybe.py
+++ b/baybe/surrogates/gaussian_process/presets/baybe.py
@@ -37,7 +37,7 @@ class _BayBENumericalKernelFactory(_SmoothedEDBONumericalKernelFactory):
     """The default numerical kernel factory for GP surrogates."""
 
 
-class BayBEKernelFactory(SmoothedEDBOKernelFactory):
+class BayBEKernelFactory(SmoothedEDBOKernelFactory):  # type: ignore[valid-type, misc]
     """The default kernel factory for GP surrogates."""
 
 

From f535cded05e6ed4ecf0dac11d127a9a2a9017fb0 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Fri, 22 May 2026 16:00:58 +0200
Subject: [PATCH 10/13] Generalize get_comp_rep_parameter_indices to accept a
 ParameterSelectorProtocol

Extends the method to accept either a parameter name (existing behavior) or
a selector, returning the combined comp-rep indices of all matching parameters.
Uses this to simplify _get_effective_dimensionality in _PureKernelFactory and
the inline dimensionality sums in the EDBO likelihood factories.
---
 baybe/searchspace/core.py                     | 59 +++++++++++++------
 .../gaussian_process/components/kernel.py     |  9 ++-
 .../gaussian_process/presets/edbo.py          |  8 +--
 .../gaussian_process/presets/edbo_smoothed.py |  8 +--
 4 files changed, 52 insertions(+), 32 deletions(-)

diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py
index 5510af704f..47e13bc4b2 100644
--- a/baybe/searchspace/core.py
+++ b/baybe/searchspace/core.py
@@ -5,7 +5,7 @@
 import gc
 from collections.abc import Iterable, Sequence
 from enum import Enum
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import pandas as pd
 from attrs import define, field
@@ -28,6 +28,9 @@
 from baybe.serialization import SerialMixin, converter, select_constructor_hook
 from baybe.utils.conversion import to_string
 
+if TYPE_CHECKING:
+    from baybe.parameters.selectors import ParameterSelectorProtocol
+
 
 class SearchSpaceType(Enum):
     """Enum class for different types of search spaces and respective compatibility."""
@@ -287,35 +290,53 @@ def n_tasks(self) -> int:
             return 1
         return len(task_param.values)
 
-    def get_comp_rep_parameter_indices(self, name: str, /) -> tuple[int, ...]:
-        """Find a parameter's column indices in the computational representation.
+    def get_comp_rep_parameter_indices(
+        self,
+        name_or_selector: str | ParameterSelectorProtocol,
+        /,
+    ) -> tuple[int, ...]:
+        """Find comp-rep column indices for a parameter selection.
+
+        When called with a parameter name, returns the indices for that single
+        parameter. When called with a
+        :class:`~baybe.parameters.selectors.ParameterSelectorProtocol`,
+        returns the combined indices for all matching parameters.
 
         Args:
-            name: The name of the parameter whose columns indices are to be retrieved.
+            name_or_selector: Either the name of a single parameter or a selector
+                that filters parameters to be included.
 
         Raises:
-            ValueError: If no parameter with the provided name exists.
-            ValueError: If more than one parameter with the provided name exists.
+            ValueError: If a parameter name is provided but no matching parameter
+                exists.
+            ValueError: If a parameter name is provided but multiple matches are
+                found.
 
         Returns:
             A tuple containing the integer indices of the columns in the computational
-            representation associated with the parameter. When the parameter is not part
-            of the computational representation, an empty tuple is returned.
+            representation associated with the selected parameter(s). When a selected
+            parameter is not part of the computational representation, it contributes
+            no indices.
         """
-        params = self.get_parameters_by_name([name])
-        if len(params) < 1:
-            raise ValueError(
-                f"There exists no parameter named '{name}' in the search space."
-            )
-        if len(params) > 1:
-            raise ValueError(
-                f"There exist multiple parameter matches for '{name}' in the search "
-                f"space."
-            )
-        p = params[0]
+        if isinstance(name_or_selector, str):
+            matched = self.get_parameters_by_name([name_or_selector])
+            if len(matched) < 1:
+                raise ValueError(
+                    f"There exists no parameter named '{name_or_selector}' in the "
+                    f"search space."
+                )
+            if len(matched) > 1:
+                raise ValueError(
+                    f"There exist multiple parameter matches for "
+                    f"'{name_or_selector}' in the search space."
+                )
+            params: list[Parameter] = list(matched)
+        else:
+            params = [p for p in self.parameters if name_or_selector(p)]
 
         return tuple(
             i
+            for p in params
             for i, col in enumerate(self.comp_rep_columns)
             if col in p.comp_rep_columns
         )
diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index ecb43f0a10..15cc19a687 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -79,11 +79,10 @@ def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...]:
 
     def _get_effective_dimensionality(self, searchspace: SearchSpace) -> int:
         """Get the number of computational columns for the selected parameters."""
-        names = self.get_parameter_names(searchspace)
-        if names is None:
-            return len(searchspace.comp_rep_columns)
-        return sum(
-            len(searchspace.get_comp_rep_parameter_indices(name)) for name in names
+        return len(
+            searchspace.get_comp_rep_parameter_indices(
+                self.parameter_selector or (lambda _: True)
+            )
         )
 
     def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None:
diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py
index 2e8596404f..c701d4dbb4 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo.py
@@ -135,10 +135,10 @@ def __call__(
         import torch
         from gpytorch.likelihoods import GaussianLikelihood
 
-        effective_dims = sum(
-            len(searchspace.get_comp_rep_parameter_indices(p.name))
-            for p in searchspace.parameters
-            if p._kind & _ParameterKind.REGULAR
+        effective_dims = len(
+            searchspace.get_comp_rep_parameter_indices(
+                lambda p: bool(p._kind & _ParameterKind.REGULAR)
+            )
         )
 
         switching_condition = _contains_encoding(
diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
index 2a9dbb441f..0ac1c6108e 100644
--- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
+++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py
@@ -110,10 +110,10 @@ def __call__(
         # 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.
         # Values outside the dimension limits will get the border value assigned.
-        effective_dims = sum(
-            len(searchspace.get_comp_rep_parameter_indices(p.name))
-            for p in searchspace.parameters
-            if p._kind & _ParameterKind.REGULAR
+        effective_dims = len(
+            searchspace.get_comp_rep_parameter_indices(
+                lambda p: bool(p._kind & _ParameterKind.REGULAR)
+            )
         )
 
         prior = GammaPrior(

From 48338191e890f93d1a10dfe9557396d358bf73e6 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Fri, 22 May 2026 16:03:34 +0200
Subject: [PATCH 11/13] Drop exception handling from
 get_comp_rep_parameter_indices

Parameter name uniqueness is already enforced as a searchspace invariant,
making the multi-match case impossible. The no-match case now returns an
empty tuple, consistent with the selector path and the existing behavior
for parameters absent from the comp-rep.
---
 baybe/searchspace/core.py | 21 +++------------------
 1 file changed, 3 insertions(+), 18 deletions(-)

diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py
index 47e13bc4b2..82273b4160 100644
--- a/baybe/searchspace/core.py
+++ b/baybe/searchspace/core.py
@@ -306,12 +306,6 @@ def get_comp_rep_parameter_indices(
             name_or_selector: Either the name of a single parameter or a selector
                 that filters parameters to be included.
 
-        Raises:
-            ValueError: If a parameter name is provided but no matching parameter
-                exists.
-            ValueError: If a parameter name is provided but multiple matches are
-                found.
-
         Returns:
             A tuple containing the integer indices of the columns in the computational
             representation associated with the selected parameter(s). When a selected
@@ -319,18 +313,9 @@ def get_comp_rep_parameter_indices(
             no indices.
         """
         if isinstance(name_or_selector, str):
-            matched = self.get_parameters_by_name([name_or_selector])
-            if len(matched) < 1:
-                raise ValueError(
-                    f"There exists no parameter named '{name_or_selector}' in the "
-                    f"search space."
-                )
-            if len(matched) > 1:
-                raise ValueError(
-                    f"There exist multiple parameter matches for "
-                    f"'{name_or_selector}' in the search space."
-                )
-            params: list[Parameter] = list(matched)
+            params: list[Parameter] = [
+                p for p in self.parameters if p.name == name_or_selector
+            ]
         else:
             params = [p for p in self.parameters if name_or_selector(p)]
 

From 4d75d5a4c258d878f717911c52d8589d2aba8580 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 22 May 2026 16:13:18 +0200
Subject: [PATCH 12/13] Add blank line for readability

Co-authored-by: Martin Fitzner <martin.fitzner@merckgroup.com>
---
 baybe/surrogates/gaussian_process/components/kernel.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index 15cc19a687..f17b289e84 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -186,6 +186,7 @@ def __call__(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor):
             self.parameter_selector = lambda p: (
                 _task_exclude_selector(p) and original_selector(p)
             )
+
         try:
             base_kernel = original_call(self, searchspace, train_x, train_y)
         finally:

From 1ccad46f65a0c5711419e63f8164b2eb9eceec69 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@gmail.com>
Date: Fri, 22 May 2026 16:40:17 +0200
Subject: [PATCH 13/13] Create sibling class instead of subclass in
 _enable_transfer_learning

---
 .../gaussian_process/components/kernel.py        | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py
index f17b289e84..2bd2bdeafe 100644
--- a/baybe/surrogates/gaussian_process/components/kernel.py
+++ b/baybe/surrogates/gaussian_process/components/kernel.py
@@ -156,12 +156,20 @@ def _enable_transfer_learning(
     # This distinction is important for serialization so that the classes can be
     # correctly identified by their names in the subclass registry
     if name is not None:
-        # Create a subclass so the original class remains unmodified.
+        # Create a sibling class so the original class remains unmodified.
+        # We use cls.__bases__ (not (cls,)) because the new class is conceptually
+        # an equivalent variant, not a specialization. Concrete (non-dunder)
+        # attributes are copied so the sibling has the same behavior.
         # __module__ must be set explicitly because the Protocol metaclass
         # would otherwise default it to "abc".
-        target_cls = type(
-            name, (cls,), {"__doc__": cls.__doc__, "__module__": cls.__module__}
-        )
+        ns = {
+            k: v
+            for k, v in cls.__dict__.items()
+            if not (k.startswith("__") and k.endswith("__"))
+        }
+        ns["__doc__"] = cls.__doc__
+        ns["__module__"] = cls.__module__
+        target_cls = type(name, cls.__bases__, ns)
     else:
         # Modify the class in-place (avoids name collision in subclass registry)
         target_cls = cls