Add TL override path to new GP factory

[kalama-ai]

Override for dispatching between PositiveIndexKernel and IndexKernel for TaskParameter

1) Optional override_transfer_learning_mode on TaskParameter

A TaskParameter can now carry an optional override_transfer_learning_mode. When left at its default (None), the kernel factory controls how the task dimension is treated. When set, it instructs the GP surrogate
to attach the requested task kernel:

# Default: factory decides (PositiveIndexKernel via BayBEKernelFactory)
task_param = TaskParameter(
    name="task",
    values=["source_1", "source_2", "target"],
    active_values=["target"],
)

# Explicit override -> PositiveIndexKernel
task_param = TaskParameter(
    name="task",
    values=["source_1", "source_2", "target"],
    active_values=["target"],
    override_transfer_learning_mode=TransferLearningMode.POSITIVE_INDEX_KERNEL,
)

# Explicit override -> IndexKernel
task_param = TaskParameter(
    name="task",
    values=["source_1", "source_2", "target"],
    active_values=["target"],
    override_transfer_learning_mode=TransferLearningMode.INDEX_KERNEL,
)

2) Dispatch in GaussianProcessSurrogate._resolve_kernel

2. Dispatch in GaussianProcessSurrogate._resolve_kernel
   a) No override — the surrogate uses the kernel factory on the full search space. Unchanged behavior.

b) Override set — the surrogate constructs the requested IndexKernel / PositiveIndexKernel on the task column and multiplies it with a task-free base kernel. How that base kernel is obtained depends on what was passed as kernel_or_factory:
Case 1 (None or BayBEKernelFactory): BayBEKernelFactory dispatches internally between _ChenNumericalKernelFactory (returns a BayBE Kernel) and _CustomScaledNumericalKernelFactory (returns a raw gpytorch kernel and calls get_comp_rep_parameter_indices). The reduced searchspace blocks get_comp_rep_parameter_indices and we can not reduce the gpytorch component to the non-task parameters. The surrogate therefore routes through ICMKernelFactory on the full search space with the prescribed task kernel directly.

A BayBEKernelFactory with a custom parameter_selector is handled like a callable factory (Case 3), i.e. evaluated on the reduced search space. This works on substance search spaces (Chen returns a name-based BayBE kernel) but currently raises IncompatibleOverrideError on non-substance ones, where the default factory yields a gpytorch kernel.

Case 2 (fixed kernel object (PlainGPComponentFactory)): PlainGPComponentFactory ignores the searchspace entirely, so calling it on a reduced space returns the same object unchanged. The task parameter is instead stripped directly from the kernel via _strip_task_from_kernel, which removes the task name from parameter_names on BasicKernels and single-level ScaleKernels, and raises for other composite kernels.

Case 3 (callable factory): The factory is called on a task-free reduced search space (SearchSpace._drop_parameters). Since the task parameter is absent, the factory produces a
task-free kernel. The result must be a BayBE Kernel. If the output is a raw gpytorch kernel it raises IncompatibleOverrideError.

In all override cases the base kernel's to_gpytorch is called on the full search space, so its active_dims align with the actual training tensor, and the task kernel is attached on the task column.

Raises IncompatibleOverrideError for: raw gpytorch kernels, composite kernels other than ScaleJKernel, task-aware factories, and factories that do not yield a BayBE kernel on the reduced space.

3) Tests

The dispatch logic is exercised directly in tests/test_kernel_factories.py, following the following scenarios

#	kernel_or_factory (with override)	Outcome
1	MaternKernel(parameter_names=("x", "Task"))	base × task (task name stripped)
2	MaternKernel(parameter_names=("x",))	base × task (already task-free)
3	factory returning a BayBE kernel on the reduced space	base × task
4	MaternKernel()	base × task (acts on all non-task parameters)
5	ScaleKernel(MaternKernel())	base × task (inner kernel stripped)
6	gpytorch.kernels.MaternKernel(active_dims=[0, 1])	IncompatibleOverrideError
7	gpytorch.kernels.MaternKernel()	IncompatibleOverrideError
8	ICMKernelFactory(...) (task-aware factory)	IncompatibleOverrideError
9	BayBEKernelFactory() (no substance parameter)	base × task (routed through ICMKernelFactory)
9b	BayBEKernelFactory(parameter_selector=...) on a non-substance space	IncompatibleOverrideError (default factory yields a gpytorch kernel)
10	IndexKernel(parameter_names=("Task",))	task kernel only (base fully stripped)
11	MaternKernel(...) * IndexKernel(...) (product kernel)	IncompatibleOverrideError

[Copilot]

Pull request overview

Adds an explicit transfer-learning override on TaskParameter and updates the GP surrogate’s kernel resolution to honor that override by dispatching between IndexKernel and PositiveIndexKernel, including new error types and test coverage updates.

Changes:

• Introduce TransferLearningMode and TaskParameter.override_transfer_learning_mode to explicitly select the task kernel type.
• Add override-aware kernel resolution in GaussianProcessSurrogate (including task-kernel attachment and incompatibility handling).
• Harden reduced-searchspace behavior with a dedicated AttributeError subclass and extend tests/strategies accordingly.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
baybe/surrogates/gaussian_process/core.py	Implements override-aware kernel dispatch/assembly and adds kernel-stripping helper.
baybe/parameters/categorical.py	Adds TaskParameter.override_transfer_learning_mode field.
baybe/parameters/enum.py	Introduces TransferLearningMode.
baybe/parameters/__init__.py	Re-exports TransferLearningMode.
baybe/searchspace/core.py	Raises UnsupportedSearchSpaceAttributeError for blocked reduced-searchspace attribute access.
baybe/exceptions.py	Adds IncompatibleOverrideError and UnsupportedSearchSpaceAttributeError.
tests/test_kernel_factories.py	Adds tests for override dispatch scenarios.
tests/test_reduced_searchspace.py	Adds test asserting the dedicated reduced-searchspace exception type.
tests/hypothesis_strategies/parameters.py	Extends TaskParameter strategy to optionally include overrides.
tests/hypothesis_strategies/kernels.py	Updates index-kernel strategy to include both index kernel variants.
CHANGELOG.md	Documents the new override feature and enum.

Comments suppressed due to low confidence (1)

baybe/surrogates/gaussian_process/core.py:447

• In the override path, when kernel_factory is a callable (i.e., not a PlainGPComponentFactory), the returned BayBE kernel is accepted as-is and later converted via to_gpytorch(searchspace=searchspace) on the full search space. If that kernel has parameter_names=None (default for many basic kernels) or otherwise still references the task parameter, it will end up acting on the task dimension as well, causing an overlap with the explicitly attached task kernel. This contradicts the intended “task-free base kernel” behavior and can yield incorrect models.

            if not isinstance(factory_kernel, Kernel):
                raise IncompatibleOverrideError(incompatible_message)
            base_spec = factory_kernel

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[AVHopp]

General logic looks good, but there is missing coverage for the case of factory being provided and one bug related to this

[Scienfitz]

not ready for review, please rebase

[AdrianSosic]

Hi @kalama-ai, thanks for taking care of this absolutely nontrivial mechanism 👍🏼