Add SoftKNNClassifierModel (#5072)

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",
 }