Respect specified model class for heterogeneous MultiTaskDatasets#5107
Open
hvarfner wants to merge 4 commits intofacebook:mainfrom
Open
Respect specified model class for heterogeneous MultiTaskDatasets#5107hvarfner wants to merge 4 commits intofacebook:mainfrom
hvarfner wants to merge 4 commits intofacebook:mainfrom
Conversation
added 4 commits
March 29, 2026 11:07
Summary: D94693361 introduced a regression: when source experiments have more parameters than target and status_quo is set, FillMissingParameters adds extra columns to target arm data during _compute_in_design. check_membership_df then returns [False] for all rows because df_cols != ss_params. All target arms are incorrectly filtered out. Fix: Override _compute_in_design in TransferLearningAdapter to use check_all_parameters_present=False. Differential Revision: D97625737
…ing (facebook#5102) Summary: When merging search spaces for transfer learning, a parameter may be FixedParameter in one experiment and ChoiceParameter in another (e.g., a parameter was fixed to a single value in the source but is tunable in the target). Previously this raised a ValueError. Now we merge them into a ChoiceParameter whose values include the union of the choice values and the fixed value. Differential Revision: D98247197
…acebook#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Wire LearnedFeatureImputation and ImputedMultiTaskGP into Ax: 1. **input_transform_argparse dispatcher**: Computes `feature_indices` and `d` from a heterogeneous MultiTaskDataset using target-first feature ordering. Validates that the dataset is a MultiTaskDataset with heterogeneous features. 2. **Storage registry**: Register ImputedMultiTaskGP in MODEL_REGISTRY and LearnedFeatureImputation in INPUT_TRANSFORM_REGISTRY. 3. **Model selection (utils.py)**: When a heterogeneous MultiTaskDataset is detected and a model class is specified (e.g. ImputedMultiTaskGP), use the specified class instead of force-overriding to HeterogeneousMTGP. Also add automatic Normalize + LearnedFeatureImputation transform chaining for ImputedMultiTaskGP. Differential Revision: D97625733
Summary: Previously, `choose_model_class` in `utils.py` would always force-override any user-specified model class to `HeterogeneousMTGP` when a heterogeneous `MultiTaskDataset` was detected. This made it impossible to use alternative models for heterogeneous transfer learning (e.g., `ImputedMultiTaskGP`). This diff changes the behavior so that when a `specified_model_class` is provided by the user (via `ModelConfig.botorch_model_class`), it is respected for heterogeneous datasets. When no model class is specified, the default behavior of selecting `HeterogeneousMTGP` is preserved. This is a prerequisite for using `ImputedMultiTaskGP` through Ax's model selection pipeline with heterogeneous search spaces. Differential Revision: D98672892
meta-cla
bot
added
the
CLA Signed
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #5107 +/- ##
==========================================
- Coverage 96.41% 96.40% -0.01%
==========================================
Files 613 613
Lines 68106 68117 +11
==========================================
+ Hits 65663 65670 +7
- Misses 2443 2447 +4 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Previously,
choose_model_classinutils.pywould always force-overrideany user-specified model class to
HeterogeneousMTGPwhen a heterogeneousMultiTaskDatasetwas detected. This made it impossible to use alternativemodels for heterogeneous transfer learning (e.g.,
ImputedMultiTaskGP).This diff changes the behavior so that when a
specified_model_classisprovided by the user (via
ModelConfig.botorch_model_class), it is respectedfor heterogeneous datasets. When no model class is specified, the default
behavior of selecting
HeterogeneousMTGPis preserved.This is a prerequisite for using
ImputedMultiTaskGPthrough Ax's modelselection pipeline with heterogeneous search spaces.
Differential Revision: D98672892