Replace GetLossClosure dispatcher with isinstance checks (#3235)

Summary:
Pull Request resolved: #3235

**Context**: See D96592835 for stack overview.

**This PR**:

This commit removes the GetLossClosure dispatcher in `get_loss_closure` and replaces it with isinstance checks for the `custom_loss` methods added in D96592835. . The routing logic is now:
1. Check for compute_loss protocol on MLL or model (added in commit 1)
2. If data_loader is provided, use _get_loss_closure_fallback_external
3. isinstance(mll, ExactMarginalLogLikelihood) -> _get_loss_closure_exact_internal
4. isinstance(mll, SumMarginalLogLikelihood) -> _get_loss_closure_sum_internal
5. Fallback -> _get_loss_closure_fallback_internal

Also cleans up:
- Removes unused type arguments (_: object, __: object, ___: None) from closure factory functions
- Removes **ignore: Any parameters that were only needed for dispatcher compatibility
- Removes **kwargs from get_loss_closure and get_loss_closure_with_grads (no callers pass extra kwargs)
- Removes imports of Dispatcher, type_bypassing_encoder, NoneType

Reviewed By: saitcakmak

Differential Revision: D96592824

fbshipit-source-id: 1a3672d9725ae3d27efb5bf78052ce18a9a5b0d7

Author: esantorella

botorch/optim/closures/model_closures.py

@@ -11,11 +11,9 @@
 from collections.abc import Callable, Sequence
 from functools import partial
 from itertools import chain, repeat
-from types import NoneType
 from typing import Any
 
 from botorch.optim.closures.core import ForwardBackwardClosure
-from botorch.utils.dispatcher import Dispatcher, type_bypassing_encoder
 from gpytorch.mlls import (
     ExactMarginalLogLikelihood,
     MarginalLogLikelihood,
@@ -24,28 +22,18 @@
 from torch import Tensor
 from torch.utils.data import DataLoader
 
-GetLossClosure = Dispatcher("get_loss_closure", encoder=type_bypassing_encoder)
-
 
 def get_loss_closure(
     mll: MarginalLogLikelihood,
     data_loader: DataLoader | None = None,
-    **kwargs: Any,
 ) -> Callable[[], Tensor]:
-    r"""Public API for GetLossClosure dispatcher.
-
-    This method, and the dispatcher that powers it, acts as a clearing house
-    for factory functions that define how ``mll`` is evaluated.
+    r"""Factory function for creating loss closures from MarginalLogLikelihoods.
 
-    Users may specify custom evaluation routines by registering a factory function
-    with GetLossClosure. These factories should be registered using the type signature
+    This method acts as a clearing house for factory functions that define how
+    ``mll`` is evaluated.
 
-        ``Type[MarginalLogLikelihood], Type[Likelihood], Type[Model],
-        Type[DataLoader]``.
-
-    The final argument, Type[DataLoader], is optional. Evaluation routines that
-    obtain training data from, e.g., ``mll.model`` should register this argument as
-    ``type(None)``.
+    Users may specify custom evaluation routines by passing an ``mll`` or an
+    ``mll.model`` with a method ``compute_custom_loss``.
 
     Args:
         mll: A MarginalLogLikelihood instance whose negative defines the loss.
@@ -61,16 +49,22 @@ def get_loss_closure(
     if hasattr(mll.model, "compute_custom_loss"):
         return partial(mll.model.compute_custom_loss, mll=mll)
 
-    return GetLossClosure(
-        mll, type(mll.likelihood), type(mll.model), data_loader, **kwargs
-    )
+    if data_loader is not None:
+        return _get_loss_closure_fallback_external(mll=mll, data_loader=data_loader)
+
+    if isinstance(mll, ExactMarginalLogLikelihood):
+        return _get_loss_closure_exact_internal(mll=mll)
+
+    if isinstance(mll, SumMarginalLogLikelihood):
+        return _get_loss_closure_sum_internal(mll=mll)
+
+    return _get_loss_closure_fallback_internal(mll=mll)
 
 
 def get_loss_closure_with_grads(
     mll: MarginalLogLikelihood,
     parameters: dict[str, Tensor],
     data_loader: DataLoader | None = None,
-    **kwargs: Any,
 ) -> ForwardBackwardClosure:
     """
     Add a backward pass to a loss closure obtained by calling
@@ -83,23 +77,18 @@ def get_loss_closure_with_grads(
         parameters: A dictionary of tensors whose ``grad`` fields are to be returned.
         data_loader: An optional DataLoader instance for cases where training
             data is passed in rather than obtained from ``mll.model``.
-        kwargs: Keyword arguments passed to ``get_loss_closure``.
 
     Returns:
         A closure that takes zero positional arguments and returns the reduced and
         negated value of ``mll`` along with the gradients of ``parameters``.
     """
-    loss_closure = get_loss_closure(mll, data_loader=data_loader, **kwargs)
+    loss_closure = get_loss_closure(mll=mll, data_loader=data_loader)
     return ForwardBackwardClosure(forward=loss_closure, parameters=parameters)
 
 
-@GetLossClosure.register(MarginalLogLikelihood, object, object, DataLoader)
 def _get_loss_closure_fallback_external(
     mll: MarginalLogLikelihood,
-    _likelihood_type: object,
-    _model_type: object,
     data_loader: DataLoader,
-    **ignore: Any,
 ) -> Callable[[], Tensor]:
     r"""Fallback loss closure with externally provided data."""
     batch_generator = chain.from_iterable(iter(data_loader) for _ in repeat(None))
@@ -120,9 +109,8 @@ def closure(**kwargs: Any) -> Tensor:
     return closure
 
 
-@GetLossClosure.register(MarginalLogLikelihood, object, object, NoneType)
 def _get_loss_closure_fallback_internal(
-    mll: MarginalLogLikelihood, _: object, __: object, ___: None, **ignore: Any
+    mll: MarginalLogLikelihood,
 ) -> Callable[[], Tensor]:
     r"""Fallback loss closure with internally managed data."""
 
@@ -134,9 +122,8 @@ def closure(**kwargs: Any) -> Tensor:
     return closure
 
 
-@GetLossClosure.register(ExactMarginalLogLikelihood, object, object, NoneType)
 def _get_loss_closure_exact_internal(
-    mll: ExactMarginalLogLikelihood, _: object, __: object, ___: None, **ignore: Any
+    mll: ExactMarginalLogLikelihood,
 ) -> Callable[[], Tensor]:
     r"""ExactMarginalLogLikelihood loss closure with internally managed data."""
 
@@ -158,9 +145,8 @@ def closure(**kwargs: Any) -> Tensor:
     return closure
 
 
-@GetLossClosure.register(SumMarginalLogLikelihood, object, object, NoneType)
 def _get_loss_closure_sum_internal(
-    mll: SumMarginalLogLikelihood, _: object, __: object, ___: None, **ignore: Any
+    mll: SumMarginalLogLikelihood,
 ) -> Callable[[], Tensor]:
     r"""SumMarginalLogLikelihood loss closure with internally managed data."""
 

test/optim/closures/test_model_closures.py

@@ -104,8 +104,8 @@ def my_compute_custom_loss(
             self.assertEqual(result(), torch.tensor(-99.0))
             del mll.model.compute_custom_loss
 
-        # Without compute_custom_loss, get_loss_closure uses the dispatcher.
-        with self.subTest("dispatcher_fallback"):
+        # Without compute_custom_loss, get_loss_closure uses isinstance checks.
+        with self.subTest("isinstance_fallback"):
             _, mlls = _get_mlls(device=self.device)
             mll = mlls[0]
             self.assertFalse(hasattr(mll, "compute_custom_loss"))