[Documentation] Incomplete documentation for implementing custom models

[slishak-PX]

The Implementing Custom Models section of the documentation makes it sound a lot more straightforward to use a custom model than it really is. So far, I've found the following extra conditions that should be mentioned (in this instance, when using qNoisyExpectedImprovement):

The posterior method must accept a posterior_transform kwarg, otherwise the following error appears:

Stack trace

File .../python3.11/site-packages/botorch/acquisition/monte_carlo.py:500, in qNoisyExpectedImprovement.__init__(self, model, X_baseline, sampler, objective, posterior_transform, X_pending, prune_baseline, cache_root, constraints, eta, marginalize_dim)
    496 CachedCholeskyMCSamplerMixin.__init__(
    497     self, model=model, cache_root=cache_root, sampler=sampler
    498 )
    499 if prune_baseline:
--> 500     X_baseline = prune_inferior_points(
    501         model=model,
    502         X=X_baseline,
    503         objective=objective,
    504         posterior_transform=posterior_transform,
    505         constraints=self._constraints,
    506         marginalize_dim=marginalize_dim,
    507     )
    508 self.register_buffer("X_baseline", X_baseline)
    509 # registering buffers for _get_samples_and_objectives in the next `if` block

File .../python3.11/site-packages/botorch/acquisition/utils.py:306, in prune_inferior_points(model, X, objective, posterior_transform, constraints, num_samples, max_frac, sampler, marginalize_dim)
    304     raise ValueError(f"max_frac must take values in (0, 1], is {max_frac}")
    305 with torch.no_grad():
--> 306     posterior = model.posterior(X=X, posterior_transform=posterior_transform)
    307 if sampler is None:
    308     sampler = get_sampler(
    309         posterior=posterior, sample_shape=torch.Size([num_samples])
    310     )

TypeError: Model.posterior() got an unexpected keyword argument 'posterior_transform'

The model additionally requires a num_outputs attribute, otherwise the following error appears (related to https://github.com/pytorch/botorch/issues/354):

Stack trace

File .../python3.11/site-packages/botorch/acquisition/monte_carlo.py:91, in MCAcquisitionFunction.__init__(self, model, sampler, objective, posterior_transform, X_pending)
     89 super().__init__(model=model)
     90 MCSamplerMixin.__init__(self, sampler=sampler)
---> 91 if objective is None and model.num_outputs != 1:
     92     if posterior_transform is None:
     93         raise UnsupportedError(
     94             "Must specify an objective or a posterior transform when using "
     95             "a multi-output model."
     96         )

File .../python3.11/site-packages/torch/nn/modules/module.py:1688, in Module.__getattr__(self, name)
   1686     if name in modules:
   1687         return modules[name]
-> 1688 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

AttributeError: 'Model' object has no attribute 'num_outputs'

A posterior sampler must be registered using @GetSampler.register(object), otherwise the following error appears:

Stack trace

File .../python3.11/site-packages/botorch/acquisition/monte_carlo.py:500, in qNoisyExpectedImprovement.__init__(self, model, X_baseline, sampler, objective, posterior_transform, X_pending, prune_baseline, cache_root, constraints, eta, marginalize_dim)
    496 CachedCholeskyMCSamplerMixin.__init__(
    497     self, model=model, cache_root=cache_root, sampler=sampler
    498 )
    499 if prune_baseline:
--> 500     X_baseline = prune_inferior_points(
    501         model=model,
    502         X=X_baseline,
    503         objective=objective,
    504         posterior_transform=posterior_transform,
    505         constraints=self._constraints,
    506         marginalize_dim=marginalize_dim,
    507     )
    508 self.register_buffer("X_baseline", X_baseline)
    509 # registering buffers for _get_samples_and_objectives in the next `if` block

File .../python3.11/site-packages/botorch/acquisition/utils.py:308, in prune_inferior_points(model, X, objective, posterior_transform, constraints, num_samples, max_frac, sampler, marginalize_dim)
    306     posterior = model.posterior(X=X, posterior_transform=posterior_transform)
    307 if sampler is None:
--> 308     sampler = get_sampler(
    309         posterior=posterior, sample_shape=torch.Size([num_samples])
    310     )
    311 samples = sampler(posterior)
    312 if objective is None:

File .../python3.11/site-packages/botorch/sampling/get_sampler.py:67, in get_sampler(posterior, sample_shape, **kwargs)
     51 r"""Get the sampler for the given posterior.
     52 
     53 The sampler can be used as `sampler(posterior)` to produce samples
   (...)
     64     The `MCSampler` object for the given posterior.
     65 """
     66 kwargs["sample_shape"] = sample_shape
---> 67 return GetSampler(posterior, **kwargs)

File .../python3.11/site-packages/botorch/utils/dispatcher.py:93, in Dispatcher.__call__(self, *args, **kwargs)
     91 func = self.__getitem__(types=types)
     92 try:
---> 93     return func(*args, **kwargs)
     94 except MDNotImplementedError:
     95     # Traverses registered methods in order, yields whenever a match is found
     96     funcs = self.dispatch_iter(*types)

File .../python3.11/site-packages/botorch/sampling/get_sampler.py:132, in _not_found_error(posterior, sample_shape, **kwargs)
    128 @GetSampler.register(object)
    129 def _not_found_error(
    130     posterior: Posterior, sample_shape: torch.Size, **kwargs: Any
    131 ) -> None:
--> 132     raise NotImplementedError(
    133         f"A registered `MCSampler` for posterior {posterior} is not found. You can "
    134         "implement and register one using `@GetSampler.register`."
    135     )

NotImplementedError: A registered `MCSampler` for posterior <__main__.ModelPosterior object at 0x7fb28dc8ed10> is not found. You can implement and register one using `@GetSampler.register`.

Furthermore, when not using stochastic sampling, the Posterior object must also implement base_sample_shape, batch_range and potentially more.

[esantorella]

Thanks for the feedback! (And for those collapsible stack traces -- very easy to navigate.) A few thoughts that might help:

• The Model base class defines the API a BoTorch model must follow. This is where you would find out what arguments the posterior method requires and that you need to implement num_outputs. Similarly for the Posterior base class and base_sample_shape and batch_range.
• The tutorial on Using a custom BoTorch. model with Ax may be helpful. Generally, people subclass some existing models such as ExactGP, which cuts down on how much work you have to do.
• That's a great point that it's not obvious that a posterior needs to be registered with GetSampler in order to be used with qNEI. You can get around this by passing a sampler argument to qNEI.

I'm also curious what your use case is and whether it might be possible to accomplish it without a custom model at all. For example, if you wanted to use a custom kernel, you could pass that as the covar_module argument to SingleTaskGP rather than defining a whole new model.

Some to-dos for improving BoTorch:

• I wonder if we should add a function like sklearn's check_estimator that programmatically lets the user know if something is not implemented right.
• We should write some clearer guidance (or at the very least not imply this is simpler than it is).
• The Model base class should have abstract methods and properties rather than methods that raise NotImplementedError so that attempting to subclass Model without defining those methods and properties causes a syntax error.
• It might be good (but maybe more effort than it's worth) to have an exception in qNEI stating that sampler cannot be None when the posterior is not registered with GetSampler.
• Finally, requirements that something be registered with a dispatcher are sneaky and not great; having get_sampler be a (class) method on a posterior might be clearer.

[slishak-PX]

Hi @esantorella,

Thanks for the response. I have subclassed Model and Posterior but at first I only implemented the abstract properties and methods, assuming from the documentation that the ones that just raised NotImplementedError would not be required, so I agree with your suggestion to change this, along with all of your other suggestions.

Specifically for qNEI and similar acquisition functions, it doesn't matter whether you provide a sampler, because the error is raised by prune_inferior_points which doesn't get passed the sampler argument. A workaround that avoids having to register the sampler is to call prune_inferior_points manually before instantiating the acquisition function. The fix that makes the most sense to me would be to pass through the sampler into prune_inferior_points when instantiating an acquisition function, but this might substantially change the behaviour of existing code.

My use case is a non-GP surrogate, hence there was not a lot of existing machinery that I could reuse. For instance, the base samples of this model are uniformly distributed rather than normally, so I had to define uniform Sobol and IID samplers (which are obviously very similar to the samplers in botorch.sampling.normal).

[esantorella]

at first I only implemented the abstract properties and methods, assuming from the documentation that the ones that just raised NotImplementedError would not be required

I looked into whether these methods could be made abstract, and unfortunately most of them can't, since there are a good number of existing subclasses that don't implement those methods. I do think the NotImplementedError API is confusing and a bad pattern, because it's hard to tell what methods need to be implemented for what purposes, but unfortunately cleaning this sort of thing up tends to be messy given the size of the codebase.

Specifically for qNEI and similar acquisition functions, it doesn't matter whether you provide a sampler, because the error is raised by prune_inferior_points which doesn't get passed the sampler argument.

Ah yes, there's even a typecheck error telling us that only a TorchPosterior works here. You might be able to get around this by passing prune_baseline=False to qNEI.

My use case is a non-GP surrogate, hence there was not a lot of existing machinery that I could reuse. For instance, the base samples of this model are uniformly distributed rather than normally, so I had to define uniform Sobol and IID samplers (which are obviously very similar to the samplers in botorch.sampling.normal).

Makes sense! Thanks.