class_docstr: **All classes in 'botorch_modular' directory are under
  construction, incomplete, and should be treated as alpha
  versions only.**
  
  Ax wrapper for BoTorch `AcquisitionFunction`, subcomponent
  of `BoTorchModel` and is not meant to be used outside of it.
  
  Args:
      surrogates: Dict of name => Surrogate model pairs, with which this acquisition
          function will be used.
      search_space_digest: A SearchSpaceDigest object containing metadata
          about the search space (e.g. bounds, parameter types).
      torch_opt_config: A TorchOptConfig object containing optimization
          arguments (e.g., objective weights, constraints).
      botorch_acqf_class: Type of BoTorch `AcquistitionFunction` that
          should be used. Subclasses of `Acquisition` often specify
          these via `default_botorch_acqf_class` attribute, in which
          case specifying one here is not required.
      options: Optional mapping of kwargs to the underlying `Acquisition
          Function` in BoTorch.
class_docstr_tok: ["*","*","All","classes","in","'botorch_modular","'","directory","are","under","construction",",","incomplete",",","and","should","be","treated","as","alpha","versions","only",".","*","*","Ax","wrapper","for","BoTorch","`","AcquisitionFunction","`",",","subcomponent","of","`","BoTorchModel","`","and","is","not","meant","to","be","used","outside","of","it",".","Args",":","surrogates",":","Dict","of","name","=",">","Surrogate","model","pairs",",","with","which","this","acquisition","function","will","be","used",".","search_space_digest",":","A","SearchSpaceDigest","object","containing","metadata","about","the","search","space","(","e.g",".","bounds",",","parameter","types",")",".","torch_opt_config",":","A","TorchOptConfig","object","containing","optimization","arguments","(","e.g.",",","objective","weights",",","constraints",")",".","botorch_acqf_class",":","Type","of","BoTorch","`","AcquistitionFunction","`","that","should","be","used",".","Subclasses","of","`","Acquisition","`","often","specify","these","via","`","default_botorch_acqf_class","`","attribute",",","in","which","case","specifying","one","here","is","not","required",".","options",":","Optional","mapping","of","kwargs","to","the","underlying","`","Acquisition","Function","`","in","BoTorch","."]
code: def __init__(
      self,
      # If using multiple Surrogates, must label primary Surrogate (typically the
      # regression Surrogate) Keys.PRIMARY_SURROGATE
      surrogates: Dict[str, Surrogate],
      search_space_digest: SearchSpaceDigest,
      torch_opt_config: TorchOptConfig,
      botorch_acqf_class: Type[AcquisitionFunction],
      options: Optional[Dict[str, Any]] = None,
  ) -> None:
      self.surrogates = surrogates
      self.options = options or {}
  
      # Compute pending and observed points for each surrogate
      Xs_pending_and_observed = {
          name: _get_X_pending_and_observed(
              Xs=surrogate.Xs,
              objective_weights=torch_opt_config.objective_weights,
              bounds=search_space_digest.bounds,
              pending_observations=torch_opt_config.pending_observations,
              outcome_constraints=torch_opt_config.outcome_constraints,
              linear_constraints=torch_opt_config.linear_constraints,
              fixed_features=torch_opt_config.fixed_features,
          )
          for name, surrogate in self.surrogates.items()
      }
  
      Xs_pending_list = [
          Xs_pending
          for Xs_pending, _ in Xs_pending_and_observed.values()
          if Xs_pending is not None
      ]
      unique_Xs_pending = (
          torch.unique(
              input=torch.cat(
                  tensors=Xs_pending_list,
                  dim=0,
              ),
              dim=0,
          )
          if len(Xs_pending_list) > 0
          else None
      )
  
      # This tensor may have some Xs that are also in pending (because they are
      # observed for some models but not others)
      Xs_observed_maybe_pending_list = [
          Xs_observed
          for _, Xs_observed in Xs_pending_and_observed.values()
          if Xs_observed is not None
      ]
      unique_Xs_observed_maybe_pending = (
          torch.unique(
              input=torch.cat(
                  tensors=Xs_observed_maybe_pending_list,
                  dim=0,
              ),
              dim=0,
          )
          if len(Xs_observed_maybe_pending_list) > 0
          else None
      )
  
      # If a point is pending on any model do not count it as observed.
      # Do this by stacking pending on top of observed, filtering repeats, then
      # removing pending points.
      # TODO[sdaulton] Is this a sound approach? Should we be doing something more
      # sophisticated here?
      if (
          unique_Xs_pending is None
          and unique_Xs_observed_maybe_pending is None
      ):
          unique_Xs_observed = None
      elif unique_Xs_pending is None:
          unique_Xs_observed = unique_Xs_observed_maybe_pending
      else:
          unique_Xs_observed = _tensor_difference(
              A=unique_Xs_pending, B=unique_Xs_observed_maybe_pending
          )
  
          if torch.numel(
              unique_Xs_observed_maybe_pending
          ) != torch.numel(unique_Xs_observed):
              logger.warning(
                  "Encountered Xs pending for some Surrogates but observed for "
                  "others. Considering these points to be pending."
              )
  
      # Store objective thresholds for all outcomes (including non-objectives).
      self._objective_thresholds: Optional[
          Tensor
      ] = torch_opt_config.objective_thresholds
      self._full_objective_weights: Tensor = (
          torch_opt_config.objective_weights
      )
      full_outcome_constraints = torch_opt_config.outcome_constraints
  
      # TODO[mpolson64] Handle more elegantly in the future. Since right now we
      # only use one objective and posterior_transform this should be fine.
      primary_surrogate = (
          self.surrogates[Keys.PRIMARY_SURROGATE]
          if len(self.surrogates) > 1
          else next(iter(self.surrogates.values()))
      )
  
      primary_Xs_pending, primary_Xs_observed = Xs_pending_and_observed[
          Keys.PRIMARY_SURROGATE
          if len(self.surrogates) > 1
          else next(iter(Xs_pending_and_observed.keys()))
      ]
  
      # Subset model only to the outcomes we need for the optimization.
      if self.options.pop(Keys.SUBSET_MODEL, True):
          subset_model_results = subset_model(
              model=primary_surrogate.model,
              objective_weights=torch_opt_config.objective_weights,
              outcome_constraints=torch_opt_config.outcome_constraints,
              objective_thresholds=torch_opt_config.objective_thresholds,
          )
          model = subset_model_results.model
          objective_weights = subset_model_results.objective_weights
          outcome_constraints = subset_model_results.outcome_constraints
          objective_thresholds = (
              subset_model_results.objective_thresholds
          )
          subset_idcs = subset_model_results.indices
      else:
          model = primary_surrogate.model
          objective_weights = torch_opt_config.objective_weights
          outcome_constraints = torch_opt_config.outcome_constraints
          objective_thresholds = torch_opt_config.objective_thresholds
          subset_idcs = None
      # If objective weights suggest multiple objectives but objective
      # thresholds are not specified, infer them using the model that
      # has already been subset to avoid re-subsetting it within
      # `inter_objective_thresholds`.
      if objective_weights.nonzero().numel() > 1 and (
          self._objective_thresholds is None
          or self._objective_thresholds[
              torch_opt_config.objective_weights != 0
          ]
          .isnan()
          .any()
      ):
          if torch_opt_config.risk_measure is not None:
              # TODO[T131759263]: modify the heuristic to support risk measures.
              raise NotImplementedError(
                  "Objective thresholds must be provided when using risk measures."
              )
          self._objective_thresholds = infer_objective_thresholds(
              model=model,
              objective_weights=self._full_objective_weights,
              outcome_constraints=full_outcome_constraints,
              X_observed=primary_Xs_observed,
              subset_idcs=subset_idcs,
              objective_thresholds=self._objective_thresholds,
          )
          objective_thresholds = (
              not_none(self._objective_thresholds)[subset_idcs]
              if subset_idcs is not None
              else self._objective_thresholds
          )
      (
          objective,
          posterior_transform,
      ) = self.get_botorch_objective_and_transform(
          botorch_acqf_class=botorch_acqf_class,
          model=model,
          objective_weights=objective_weights,
          objective_thresholds=objective_thresholds,
          outcome_constraints=outcome_constraints,
          X_observed=primary_Xs_observed,
          risk_measure=torch_opt_config.risk_measure,
      )
  
      model_deps = self.compute_model_dependencies(
          surrogates=self.surrogates,
          search_space_digest=search_space_digest,
          torch_opt_config=dataclasses.replace(
              torch_opt_config,
              objective_weights=objective_weights,
              outcome_constraints=outcome_constraints,
              objective_thresholds=objective_thresholds,
          ),
          options=self.options,
      )
  
      acqf_model_kwarg = (
          {
              "model_dict": ModelDict(
                  **{
                      name: surrogate.model
                      for name, surrogate in self.surrogates.items()
                  }
              )
          }
          if len(self.surrogates) > 1
          else {"model": model}
      )
      target_fidelities = {
          k: v
          for k, v in search_space_digest.target_values.items()
          if k in search_space_digest.fidelity_features
      }
      input_constructor_kwargs = {
          "X_baseline": unique_Xs_observed,
          "X_pending": unique_Xs_pending,
          "objective_thresholds": objective_thresholds,
          "constraints": get_outcome_constraint_transforms(
              outcome_constraints=outcome_constraints
          ),
          "target_fidelities": target_fidelities,
          "bounds": search_space_digest.bounds,
          **acqf_model_kwarg,
          **model_deps,
          **self.options,
      }
      input_constructor = get_acqf_input_constructor(botorch_acqf_class)
      # Handle multi-dataset surrogates - TODO: Improve this
      # If there is only one SupervisedDataset return it alone
      if (
          len(self.surrogates) == 1
          and len(next(iter(self.surrogates.values())).training_data)
          == 1
      ):
          training_data = next(
              iter(self.surrogates.values())
          ).training_data[0]
      else:
          tdicts = (
              dict(
                  zip(
                      not_none(surrogate._outcomes),
                      surrogate.training_data,
                  )
              )
              for surrogate in self.surrogates.values()
          )
          # outcome_name => Dataset
          training_data = functools.reduce(
              lambda x, y: {**x, **y}, tdicts
          )
  
      acqf_inputs = input_constructor(
          training_data=training_data,
          objective=objective,
          posterior_transform=posterior_transform,
          **input_constructor_kwargs,
      )
      self.acqf = botorch_acqf_class(**acqf_inputs)  # pyre-ignore [45]
      self.X_pending: Optional[Tensor] = unique_Xs_pending
      self.X_observed: Tensor = not_none(unique_Xs_observed)
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  import operator
  import warnings
  from functools import partial, reduce
  from itertools import product
  from logging import Logger
  from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Type
  import torch
  from ax.core.search_space import SearchSpaceDigest
  from ax.exceptions.core import AxWarning, SearchSpaceExhausted
  from ax.models.model_utils import enumerate_discrete_combinations, mk_discrete_choices
  from ax.models.torch.botorch_modular.optimizer_argparse import optimizer_argparse
  from ax.models.torch.botorch_modular.surrogate import Surrogate
  from ax.models.torch.botorch_modular.utils import _tensor_difference, get_post_processing_func
  from ax.models.torch.botorch_moo_defaults import infer_objective_thresholds
  from ax.models.torch.utils import _get_X_pending_and_observed, get_botorch_objective_and_transform, subset_model
  from ax.models.torch_base import TorchOptConfig
  from ax.utils.common.base import Base
  from ax.utils.common.constants import Keys
  from ax.utils.common.logger import get_logger
  from ax.utils.common.typeutils import not_none
  from botorch.acquisition.acquisition import AcquisitionFunction
  from botorch.acquisition.input_constructors import get_acqf_input_constructor
  from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
  from botorch.acquisition.objective import MCAcquisitionObjective, PosteriorTransform
  from botorch.acquisition.risk_measures import RiskMeasureMCObjective
  from botorch.models.model import Model, ModelDict
  from botorch.optim.optimize import optimize_acqf, optimize_acqf_discrete, optimize_acqf_discrete_local_search, optimize_acqf_mixed
  from botorch.utils.constraints import get_outcome_constraint_transforms
  from torch import Tensor
label_desc: Use image node_id 1 to create a new Acquisition object from inherited base classes: Base with example: obj = Acquisition(surrogates, search_space_digest, torch_opt_config, botorch_acqf_class, options)