acquisition function wrapper (meta-pytorch#1532)

Summary:
Pull Request resolved: meta-pytorch#1532

Add a wrapper for modifying inputs/outputs. This is useful for not only probabilistic reparameterization, but will also simplify other integrated AFs (e.g. MCMC) as well as fixed feature AFs and things like prior-guided AFs

Differential Revision: D41629186

fbshipit-source-id: 6dcfdb4ebf4dd316f361d1d3b0c9bfa1d54ff2d1

Author: sdaulton

botorch/acquisition/wrapper.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+r"""
+A wrapper classes around AcquisitionFunctions to modify inputs and outputs.
+"""
+
+from __future__ import annotations
+
+from botorch.acquisition.acquisition import AcquisitionFunction
+from torch.nn import Module
+
+
+class AcquisitionFunctionWrapper(AcquisitionFunction):
+    r"""Abstract acquisition wrapper."""
+
+    def __init__(self, acq_function: AcquisitionFunction) -> None:
+        Module.__init__(self)
+        self.__class__ = type(
+            acq_function.__class__.__name__,
+            (self.__class__, acq_function.__class__),
+            {},
+        )
+        self.__dict__ = acq_function.__dict__
+        self.acq_function = acq_function