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.

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)
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

[AVHopp]

Should this docstring maybe already contain some more information about how conflicts are handled?

[AVHopp]

In general - where and how do we best document this feature/behavior? ALso in the docstring of the GPSurrogate or its fields?

[kalama-ai]

I added the main description to the TaskParameter itsefl and a brief comment to the surrogate's field: c7862ad. Are you fine with this version?

[AVHopp]

It feels weird to me that this function is a private helper within Gaussian_Process/core.py. Shouldn't it eiter live directly with kernels or at some utility function for kernels?

[kalama-ai]

I moved the helper to kernels: ae3c5c5. I introduced a new method for kernels _without_parameter that will return the same kernel operating on the reduced set of parameters for all kernels for which we can easily define the reduces version. Other kernels, like additive ones will raise a TypeError. Does this work for you?

[AVHopp]

Aren't we missing a test case for having a factory here? I only see None and explicit kernels

[kalama-ai]

I added a test for factories returning gpytorch kernels and task-free BayBE kernels here: 42fc1c7. For factories returning task-aware BayBE factories see your other comment.

[kalama-ai]

Added a test for factories returning task-aware BayBE factories here: bfa8983.

[AVHopp]

This field is ignored for all other models that are not GPSurrogates, should this be flagged/mentioned somewhere/raise an error/warning?

[kalama-ai]

Good point. The thing is that currently GPSurrogate is the only surrogate that has supports_transfer_learning=True and can work with a TaskParameter. Any other surrogate will already raise an ValueError if combined with a TaskParameter and the override field will never be ignored. If we ever add another surrogate for TL that can not use the override we should definitely coem back to this.

[AVHopp]

There is an issue here in the case of a factory which produces a kernel that has parameter_names=None: In this case, we land in the "else" branch above, and factory_kernel is then a kernel built on the reduced search space, but with parameter_names=None, so it is unclear that this should be "task-free". Thus, the moment we set base_kernel, this kernel spans all columns of the search space, including the task-related one. But we still multiply the task kernel, hence modelling the task twice.

[AVHopp]

Code demonstrating this:

def _unnamed_kernel_factory(searchspace, objective, measurements):
    """A callable factory returning a BayBE kernel with `parameter_names=None`."""
    return MaternKernel()  # parameter_names defaults to None -> "acts on all dims"


def test_pathD_double_covers_task_column(monkeypatch):
    """A callable factory returning an unnamed kernel double-covers the task column."""
    monkeypatch.setenv("BAYBE_DISABLE_CUSTOM_KERNEL_WARNING", "True")

    task = TaskParameter(
        "Task",
        ["A", "B", "C"],
        active_values=["A"],
        override_transfer_learning_mode=TransferLearningMode.INDEX_KERNEL,
    )
    num = NumericalDiscreteParameter("x", [1, 2, 3, 4, 5])
    searchspace = SearchSpace.from_product([num, task])
    context = _ModelContext(
        searchspace, NumericalTarget("y").to_objective(), pd.DataFrame()
    )
    task_idx = context.task_idx

    surrogate = GaussianProcessSurrogate(kernel_or_factory=_unnamed_kernel_factory)
    kernel = surrogate._resolve_kernel(context)

    base_kernel, task_kernel = kernel.kernels

    print(f"\ntask column index         = {task_idx}")
    print(f"base_kernel.active_dims   = {base_kernel.active_dims}")
    print(f"task_kernel.active_dims   = {task_kernel.active_dims}")

    # The task kernel correctly acts only on the task column.
    assert set(task_kernel.active_dims.tolist()) == {task_idx}

    # In gpytorch, `active_dims is None` means "acts on every input dimension" -
    # i.e. the base kernel also covers the task column, modeling it twice. A
    # properly task-free base kernel must instead be restricted to a subset.
    assert base_kernel.active_dims is not None

When this is added as a test, it fails

[kalama-ai]

Added a fix here: bfa8983. A factory with parameter_names=None is now reduced to the non-task paremeters. Included a test.

[kalama-ai]

Good catch btw :) Thanks!

[Scienfitz]

not ready for review, please rebase

[Scienfitz]

could you amend and squash some commits such that DISPATCHING_PLAN does not appear anymore at all in this PR?

[kalama-ai]

Will squash all commits in the end and update commit message.

[Scienfitz]

please make your agent respect AGENTS.md more, in particular the instruction to not have blown-up commit bodies with redundant information

Example: Rename X into Y, this does not have to be written in the commit body, its apparent in the code and provides no additional info

[kalama-ai]

Will squash all commits in the end and update commit message.

[Scienfitz]

missing guards for non_task_names, eg could be empty

is there a slightly different design possible that wouldnt require to provide both sets of names?

[Scienfitz]

weve tests covering the construction, but could it also be worth to add 1 simple test that uses an override in test_iterations?

[Scienfitz]

im not entirely sure myself, but what about the parameter_selector the baybekernelfactory has? do we need to pass that on or ignore it like done here?

[Scienfitz]

usually would prefer isinstance or is there a reason you chose this? (derived classes would fail this check)

[Scienfitz]

does this turn intoa. proper reference? or do we need somtehing else like eg " :class:.TaskParameter"

[Scienfitz]

also here isinstance is preferred over name tests