Add SoftKNNClassifierModel (facebook#5072)

Summary:

X-link: meta-pytorch/botorch#3243

Add a differentiable Soft K-Nearest Neighbors classifier model for failure-aware
Bayesian optimization. Unlike tree-based classifiers (RF, XGBoost), SoftKNN is
fully differentiable, enabling gradient-based acquisition function optimization.

The model uses Gaussian kernel weights:
  P(y=1|x) = sum(w_i * y_i) / sum(w_i)
  where w_i = exp(-||x - x_i||^2 / (2 * sigma^2))

Implements construct_inputs classmethod for seamless Ax integration.

Differential Revision: D90894389

Author: Sunny Shen

ax/generators/torch/botorch_modular/utils.py

@@ -37,6 +37,7 @@
 from botorch.exceptions.errors import BotorchError, CandidateGenerationError
 from botorch.fit import fit_fully_bayesian_model_nuts, fit_gpytorch_mll
 from botorch.models import PairwiseLaplaceMarginalLogLikelihood
+from botorch.models.classifier import SoftKNNClassifierModel
 from botorch.models.fully_bayesian import AbstractFullyBayesianSingleTaskGP
 from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 from botorch.models.gp_regression import SingleTaskGP
@@ -808,6 +809,15 @@ def _fit_botorch_model_fully_bayesian_nuts(
     fit_fully_bayesian_model_nuts(model, **mll_options)
 
 
+@fit_botorch_model.register(SoftKNNClassifierModel)
+def _fit_botorch_model_classifier(
+    model: SoftKNNClassifierModel,
+    mll_class: type[MarginalLogLikelihood],
+    mll_options: dict[str, Any] | None = None,
+) -> None:
+    """Classifier models fit themselves in __init__(), so no-op here."""
+
+
 @fit_botorch_model.register(object)
 def _fit_botorch_model_not_implemented(
     model: Model,

ax/storage/botorch_modular_registry.py

@@ -60,6 +60,7 @@
     AnalyticExpectedUtilityOfBestOption,
     qExpectedUtilityOfBestOption,
 )
+from botorch.models.classifier import SoftKNNClassifierModel
 from botorch.models.contextual import LCEAGP
 from botorch.models.fully_bayesian import (
     FullyBayesianLinearSingleTaskGP,
@@ -145,6 +146,8 @@
     AdditiveMapSaasSingleTaskGP: "AdditiveMapSaasSingleTaskGP",
     EnsembleMapSaasSingleTaskGP: "EnsembleMapSaasSingleTaskGP",
     HeterogeneousMTGP: "HeterogeneousMTGP",
+    # Classifier models
+    SoftKNNClassifierModel: "SoftKNNClassifierModel",
 }