Transfer Learning Decorator

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
@@ -119,6 +120,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."""
@@ -154,18 +196,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__(

baybe/surrogates/gaussian_process/presets/baybe.py

@@ -22,38 +22,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

baybe/surrogates/gaussian_process/presets/chen.py

@@ -6,19 +6,14 @@
 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.kernel import (
+    _enable_transfer_learning,
     _PureKernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
@@ -33,19 +28,14 @@
     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
 
     _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

baybe/surrogates/gaussian_process/presets/edbo.py

@@ -6,22 +6,18 @@
 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
 from baybe.surrogates.gaussian_process.components.kernel import (
+    _enable_transfer_learning,
     _PureKernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
@@ -56,6 +52,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`.
@@ -67,12 +64,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

baybe/surrogates/gaussian_process/presets/edbo_smoothed.py

@@ -6,19 +6,15 @@
 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.kernel import (
+    _enable_transfer_learning,
     _PureKernelFactory,
 )
 from baybe.surrogates.gaussian_process.components.likelihood import (
@@ -39,23 +35,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
@@ -88,6 +73,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."""
 

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",