Add input_transform_argparse dispatcher for LearnedFeatureImputation (#5106) by hvarfner · Pull Request #5106 · meta/Ax

hvarfner · 2026-03-29T17:10:56Z

Summary:
X-link: https://github.com/facebookexternal/botorch_fb/pull/34

Wire LearnedFeatureImputation and ImputedMultiTaskGP into Ax:

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.
Storage registry: Register ImputedMultiTaskGP in MODEL_REGISTRY and
LearnedFeatureImputation in INPUT_TRANSFORM_REGISTRY.
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.

Reviewed By: sdaulton

Differential Revision: D97625733

meta-codesync · 2026-03-29T17:11:16Z

@hvarfner has exported this pull request. If you are a Meta employee, you can view the originating Diff in D97625733.

…eta#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Pull Request resolved: meta#5106 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

codecov-commenter · 2026-03-29T17:50:40Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.43%. Comparing base (98c7fd1) to head (aea2170).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #5106   +/-   ##
=======================================
  Coverage   96.43%   96.43%           
=======================================
  Files         617      617           
  Lines       69309    69379   +70     
=======================================
+ Hits        66837    66907   +70     
  Misses       2472     2472

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…eta#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

…eta#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Pull Request resolved: meta#5106 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

…eta#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

…eta#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Pull Request resolved: meta#5106 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

…eta#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

…eta#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Pull Request resolved: meta#5106 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

…eta#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

…eta#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Pull Request resolved: meta#5106 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

…eta#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. Reviewed By: sdaulton Differential Revision: D97625733

…eta#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Pull Request resolved: meta#5106 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

…eta#5106) Summary: X-link: https://github.com/facebookexternal/botorch_fb/pull/34 Pull Request resolved: meta#5106 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. Reviewed By: sdaulton Differential Revision: D97625733

…eta#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. Reviewed By: sdaulton Differential Revision: D97625733

meta-codesync · 2026-04-24T03:58:11Z

This pull request has been merged in eeb78fe.

hvarfner force-pushed the export-D97625733 branch from a629e6f to 14dec0a Compare March 29, 2026 17:11

meta-cla Bot added the CLA Signed Do not delete this pull request or issue due to inactivity. label Mar 29, 2026

meta-codesync Bot added fb-exported meta-exported labels Mar 29, 2026

hvarfner force-pushed the export-D97625733 branch from 14dec0a to 45f080a Compare March 29, 2026 17:14

meta-codesync Bot changed the title ~~Add input_transform_argparse dispatcher for LearnedFeatureImputation~~ Add input_transform_argparse dispatcher for LearnedFeatureImputation (#5106) Mar 29, 2026

hvarfner force-pushed the export-D97625733 branch from 45f080a to 397aeaa Compare April 21, 2026 15:24

hvarfner force-pushed the export-D97625733 branch from 397aeaa to d33f053 Compare April 21, 2026 16:30

hvarfner force-pushed the export-D97625733 branch from d33f053 to 05c67bf Compare April 21, 2026 19:50

hvarfner force-pushed the export-D97625733 branch from 05c67bf to f599cc2 Compare April 21, 2026 23:06

hvarfner force-pushed the export-D97625733 branch from f599cc2 to c5d35f1 Compare April 21, 2026 23:23

hvarfner force-pushed the export-D97625733 branch 3 times, most recently from a1d0253 to 2ec2c47 Compare April 23, 2026 16:00

hvarfner force-pushed the export-D97625733 branch from 2ec2c47 to 252b618 Compare April 23, 2026 19:28

hvarfner force-pushed the export-D97625733 branch from 252b618 to b180c3d Compare April 23, 2026 19:32

hvarfner force-pushed the export-D97625733 branch from b180c3d to aea2170 Compare April 24, 2026 02:59

meta-codesync Bot closed this in eeb78fe Apr 24, 2026

facebook-github-tools Bot added the Merged label Apr 24, 2026

Navigation Menu

Repository navigation

Add input_transform_argparse dispatcher for LearnedFeatureImputation (#5106) - #5106

Add input_transform_argparse dispatcher for LearnedFeatureImputation (#5106)#5106
hvarfner wants to merge 1 commit into
meta:mainfrom
hvarfner:export-D97625733

hvarfner commented Mar 29, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

meta-codesync Bot commented Mar 29, 2026

Uh oh!

codecov-commenter commented Mar 29, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented Apr 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

hvarfner commented Mar 29, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented Mar 29, 2026

Uh oh!

codecov-commenter commented Mar 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

meta-codesync Bot commented Apr 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

hvarfner commented Mar 29, 2026 •

edited by meta-codesync Bot

Loading

codecov-commenter commented Mar 29, 2026 •

edited

Loading