Add input_transform_argparse dispatcher for LearnedFeatureImputation (#5106)

[hvarfner]

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]

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

[codecov-commenter]

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.

[meta-codesync]

This pull request has been merged in eeb78fe.