ax.models¶
Base Models¶
ax.models.base¶

class
ax.models.base.
Model
[source]¶ Bases:
object
Base class for an Ax model.
Note: the core methods each model has: fit, predict, gen, cross_validate, and best_point are not present in this base class, because the signatures for those methods vary based on the type of the model. This class only contains the methods that all models have in common and for which they all share the signature.
ax.models.discrete_base module¶

class
ax.models.discrete_base.
DiscreteModel
[source]¶ Bases:
ax.models.base.Model
This class specifies the interface for a model based on discrete parameters.
These methods should be implemented to have access to all of the features of Ax.

best_point
(n: int, parameter_values: List[List[Optional[Union[str, bool, float, int]]]], objective_weights: Optional[numpy.ndarray], outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, Optional[Union[str, bool, float, int]]]] = None, pending_observations: Optional[List[List[List[Optional[Union[str, bool, float, int]]]]]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Optional[List[Optional[Union[str, bool, float, int]]]][source]¶ Obtains the point that has the best value according to the model prediction and its model predictions.
 Returns
(1 x d) parameter value list representing the point with the best value according to the model prediction. None if this function is not implemented for the given model.

cross_validate
(Xs_train: List[List[List[Optional[Union[str, bool, float, int]]]]], Ys_train: List[List[float]], Yvars_train: List[List[float]], X_test: List[List[Optional[Union[str, bool, float, int]]]]) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m lists X of parameterizations (each parameterization is a list of parameter values of length d), each of length k_i, for each outcome.
Ys_train – The corresponding list of m lists Y, each of length k_i, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – List of the j parameterizations at which to make predictions.
 Returns
2element tuple containing
(j x m) array of outcome predictions at X.
(j x m x m) array of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

fit
(Xs: List[List[List[Optional[Union[str, bool, float, int]]]]], Ys: List[List[float]], Yvars: List[List[float]], parameter_values: List[List[Optional[Union[str, bool, float, int]]]], outcome_names: List[str]) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m lists X of parameterizations (each parameterization is a list of parameter values of length d), each of length k_i, for each outcome.
Ys – The corresponding list of m lists Y, each of length k_i, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
parameter_values – A list of possible values for each parameter.
outcome_names – A list of m outcome names.

gen
(n: int, parameter_values: List[List[Optional[Union[str, bool, float, int]]]], objective_weights: Optional[numpy.ndarray], outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, Optional[Union[str, bool, float, int]]]] = None, pending_observations: Optional[List[List[List[Optional[Union[str, bool, float, int]]]]]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Tuple[List[List[Optional[Union[str, bool, float, int]]]], List[float], Dict[str, Any]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
parameter_values – A list of possible values for each parameter.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m lists of parameterizations (each parameterization is a list of parameter values of length d), each of length k_i, for each outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
 Returns
2element tuple containing
List of n generated points, where each point is represented by a list of parameter values.
List of weights for each of the n points.

predict
(X: List[List[Optional[Union[str, bool, float, int]]]]) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Predict
 Parameters
X – List of the j parameterizations at which to make predictions.
 Returns
2element tuple containing
(j x m) array of outcome predictions at X.
(j x m x m) array of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

ax.models.model_utils module¶

ax.models.model_utils.
add_fixed_features
(tunable_points: numpy.ndarray, d: int, fixed_features: Optional[Dict[int, float]], tunable_feature_indices: numpy.ndarray) → numpy.ndarray[source]¶ Add fixed features to points in tunable space.
 Parameters
tunable_points – Points in tunable space.
d – Dimension of parameter space.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
tunable_feature_indices – Parameter indices (in d) which are tunable.
 Returns
Points in the full ddimensional space, defined by bounds.
 Return type
points

ax.models.model_utils.
as_array
(x: Union[torch.Tensor, numpy.ndarray, Tuple[Union[torch.Tensor, numpy.ndarray], …]]) → Union[numpy.ndarray, Tuple[numpy.ndarray, …]][source]¶ Convert every item in a tuple of tensors/arrays into an array.
 Parameters
x – A tensor, array, or a tuple of potentially mixed tensors and arrays.
 Returns
x, with everything converted to array.

ax.models.model_utils.
best_in_sample_point
(Xs: Union[List[torch.Tensor], List[numpy.ndarray]], model: Union[ax.models.numpy_base.NumpyModel, ax.models.torch_base.TorchModel], bounds: List[Tuple[float, float]], objective_weights: Optional[Union[torch.Tensor, numpy.ndarray]], outcome_constraints: Optional[Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, linear_constraints: Optional[Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Optional[Tuple[Union[torch.Tensor, numpy.ndarray], float]][source]¶ Select the best point that has been observed.
Implements two approaches to selecting the best point.
For both approaches, only points that satisfy parameter space constraints (bounds, linear_constraints, fixed_features) will be returned. Points must also be observed for all objective and constraint outcomes. Returned points may violate outcome constraints, depending on the method below.
1: Select the point that maximizes the expected utility (objective_weights^T posterior_objective_means  baseline) * Prob(feasible) Here baseline should be selected so that at least one point has positive utility. It can be specified in the options dict, otherwise min (objective_weights^T posterior_objective_means) will be used, where the min is over observed points.
2: Select the bestobjective point that is feasible with at least probability p.
The following quantities may be specified in the options dict:
best_point_method: ‘max_utility’ (default) or ‘feasible_threshold’ to select between the two approaches described above.
utility_baseline: Value for the baseline used in max_utility approach. If not provided, defaults to min objective value.
probability_threshold: Threshold for the feasible_threshold approach. Defaults to p=0.95.
feasibility_mc_samples: Number of MC samples used for estimating the probability of feasibility (defaults 10k).
 Parameters
Xs – Training data for the points, among which to select the best.
model – Numpy or Torch model.
bounds – A list of (lower, upper) tuples for each feature.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
options – A config dictionary with settings described above.
 Returns
darray of the best point,
utility at the best point.
 Return type
A twoelement tuple or None if no feasible point exist. In tuple

ax.models.model_utils.
best_observed_point
(model: Union[ax.models.numpy_base.NumpyModel, ax.models.torch_base.TorchModel], bounds: List[Tuple[float, float]], objective_weights: Optional[Union[torch.Tensor, numpy.ndarray]], outcome_constraints: Optional[Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, linear_constraints: Optional[Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Optional[Union[torch.Tensor, numpy.ndarray]][source]¶ Select the best point that has been observed.
Implements two approaches to selecting the best point.
For both approaches, only points that satisfy parameter space constraints (bounds, linear_constraints, fixed_features) will be returned. Points must also be observed for all objective and constraint outcomes. Returned points may violate outcome constraints, depending on the method below.
1: Select the point that maximizes the expected utility (objective_weights^T posterior_objective_means  baseline) * Prob(feasible) Here baseline should be selected so that at least one point has positive utility. It can be specified in the options dict, otherwise min (objective_weights^T posterior_objective_means) will be used, where the min is over observed points.
2: Select the bestobjective point that is feasible with at least probability p.
The following quantities may be specified in the options dict:
best_point_method: ‘max_utility’ (default) or ‘feasible_threshold’ to select between the two approaches described above.
utility_baseline: Value for the baseline used in max_utility approach. If not provided, defaults to min objective value.
probability_threshold: Threshold for the feasible_threshold approach. Defaults to p=0.95.
feasibility_mc_samples: Number of MC samples used for estimating the probability of feasibility (defaults 10k).
 Parameters
model – Numpy or Torch model.
bounds – A list of (lower, upper) tuples for each feature.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
options – A config dictionary with settings described above.
 Returns
A darray of the best point, or None if no feasible point exists.

ax.models.model_utils.
check_duplicate
(point: numpy.ndarray, points: numpy.ndarray) → bool[source]¶ Check if a point exists in another array.
 Parameters
point – Newly generated point to check.
points – Points previously generated.
 Returns
True if the point is contained in points, else False

ax.models.model_utils.
check_param_constraints
(linear_constraints: Tuple[numpy.ndarray, numpy.ndarray], point: numpy.ndarray) → Tuple[bool, numpy.ndarray][source]¶ Check if a point satisfies parameter constraints.
 Parameters
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
point – A candidate point in ddimensional space, as a (1 x d) matrix.
 Returns
2element tuple containing
Flag that is True if all constraints are satisfied by the point.
Indices of constraints which are violated by the point.

ax.models.model_utils.
enumerate_discrete_combinations
(discrete_choices: Dict[int, List[Union[int, float]]]) → List[Dict[int, Union[int, float]]][source]¶

ax.models.model_utils.
filter_constraints_and_fixed_features
(X: Union[torch.Tensor, numpy.ndarray], bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, fixed_features: Optional[Dict[int, float]] = None) → Union[torch.Tensor, numpy.ndarray][source]¶ Filter points to those that satisfy bounds, linear_constraints, and fixed_features.
 Parameters
X – An tensor or array of points.
bounds – A list of (lower, upper) tuples for each feature.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
 Returns
Feasible points.

ax.models.model_utils.
get_observed
(Xs: Union[List[torch.Tensor], List[numpy.ndarray]], objective_weights: Union[torch.Tensor, numpy.ndarray], outcome_constraints: Optional[Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None) → Union[torch.Tensor, numpy.ndarray][source]¶ Filter points to those that are observed for objective outcomes and outcomes that show up in outcome_constraints (if there are any).
 Parameters
Xs – A list of m (k_i x d) feature matrices X. Number of rows k_i can vary from i=1,…,m.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
 Returns
Points observed for all objective outcomes and outcome constraints.

ax.models.model_utils.
mk_discrete_choices
(ssd: ax.core.search_space.SearchSpaceDigest, fixed_features: Optional[Dict[int, float]] = None) → Dict[int, List[Union[int, float]]][source]¶

ax.models.model_utils.
rejection_sample
(gen_unconstrained: Callable[[int, int, numpy.ndarray, Optional[Dict[int, float]]], numpy.ndarray], n: int, d: int, tunable_feature_indices: numpy.ndarray, linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, deduplicate: bool = False, max_draws: Optional[int] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None, existing_points: Optional[numpy.ndarray] = None) → Tuple[numpy.ndarray, int][source]¶ Rejection sample in parameter space.
Models must implement a gen_unconstrained method in order to support rejection sampling via this utility.

ax.models.model_utils.
tunable_feature_indices
(bounds: List[Tuple[float, float]], fixed_features: Optional[Dict[int, float]] = None) → numpy.ndarray[source]¶ Get the feature indices of tunable features.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
 Returns
The indices of tunable features.

ax.models.model_utils.
validate_bounds
(bounds: List[Tuple[float, float]], fixed_feature_indices: numpy.ndarray) → None[source]¶ Ensure the requested space is [0,1]^d.
 Parameters
bounds – A list of d (lower, upper) tuples for each column of X.
fixed_feature_indices – Indices of features which are fixed at a particular value.
ax.models.numpy_base module¶

class
ax.models.numpy_base.
NumpyModel
[source]¶ Bases:
ax.models.base.Model
This class specifies the interface for a numpybased model.
These methods should be implemented to have access to all of the features of Ax.

best_point
(bounds: List[Tuple[float, float]], objective_weights: numpy.ndarray, outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Optional[numpy.ndarray][source]¶ Identify the current best point, satisfying the constraints in the same format as to gen.
Return None if no such point can be identified.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
 Returns
A darray of the best point.

cross_validate
(Xs_train: List[numpy.ndarray], Ys_train: List[numpy.ndarray], Yvars_train: List[numpy.ndarray], X_test: numpy.ndarray) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m (k_i x d) feature matrices X. Number of rows k_i can vary from i=1,…,m.
Ys_train – The corresponding list of m (k_i x 1) outcome arrays Y, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – (j x d) array of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) array of outcome predictions at X.
(j x m x m) array of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

fit
(Xs: List[numpy.ndarray], Ys: List[numpy.ndarray], Yvars: List[numpy.ndarray], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature matrices X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome arrays Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

gen
(n: int, bounds: List[Tuple[float, float]], objective_weights: numpy.ndarray, outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[numpy.ndarray]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None) → Tuple[numpy.ndarray, numpy.ndarray, Dict[str, Any], Optional[List[Optional[Dict[str, Any]]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature arrays X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result (xbest) appropriately (i.e., according to roundtrip transformations)
 Returns
4element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
Generation metadata
 Dictionary of modelspecific metadata for the given
generation candidates

predict
(X: numpy.ndarray) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Predict
 Parameters
X – (j x d) array of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) array of outcome predictions at X.
(j x m x m) array of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

update
(Xs: List[numpy.ndarray], Ys: List[numpy.ndarray], Yvars: List[numpy.ndarray], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, **kwargs: Any) → None[source]¶ Update the model.
Updating the model requires both existing and additional data. The data passed into this method will become the new training data.
 Parameters
Xs – Existing + additional data for the model, in the same format as for fit.
Ys – Existing + additional data for the model, in the same format as for fit.
Yvars – Existing + additional data for the model, in the same format as for fit.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

ax.models.torch_base module¶

class
ax.models.torch_base.
TorchModel
[source]¶ Bases:
ax.models.base.Model
This class specifies the interface for a torchbased model.
These methods should be implemented to have access to all of the features of Ax.

best_point
(bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ Identify the current best point, satisfying the constraints in the same format as to gen.
Return None if no such point can be identified.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
dtensor of the best point.

cross_validate
(Xs_train: List[torch.Tensor], Ys_train: List[torch.Tensor], Yvars_train: List[torch.Tensor], X_test: torch.Tensor, search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str]) → Tuple[torch.Tensor, torch.Tensor][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys_train – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – (j x d) tensor of the j points at which to make predictions.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

evaluate_acquisition_function
(X: torch.Tensor, search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, objective_thresholds: Optional[torch.Tensor] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, acq_options: Optional[Dict[str, Any]] = None) → torch.Tensor[source]¶ Evaluate the acquisition function on the candidate set X.
 Parameters
X – (j x d) tensor of the j points at which to evaluate the acquisition function.
search_space_digest – A dataclass used to compactly represent a search space.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
objective_thresholds – The mdim tensor of objective thresholds. There is one for each modeled metric.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be held fixed during the evaluation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
acq_options – Keyword arguments used to contruct the acquisition function.
 Returns
A singleelement tensor with the acquisition value for these points.

fit
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

gen
(n: int, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], Optional[List[Optional[Dict[str, Any]]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
4element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
Generation metadata
 Dictionary of modelspecific metadata for the given
generation candidates

predict
(X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predict
 Parameters
X – (j x d) tensor of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

update
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Update the model.
Updating the model requires both existing and additional data. The data passed into this method will become the new training data.
 Parameters
Xs – Existing + additional data for the model, in the same format as for fit.
Ys – Existing + additional data for the model, in the same format as for fit.
Yvars – Existing + additional data for the model, in the same format as for fit.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

Discrete Models¶
ax.models.discrete.eb_thompson module¶

class
ax.models.discrete.eb_thompson.
EmpiricalBayesThompsonSampler
(num_samples: int = 10000, min_weight: Optional[float] = None, uniform_weights: bool = False)[source]¶ Bases:
ax.models.discrete.thompson.ThompsonSampler
Generator for Thompson sampling using Empirical Bayes estimates.
The generator applies positivepart JamesStein Estimator to the data passed in via fit and then performs Thompson Sampling.
ax.models.discrete.full_factorial module¶

class
ax.models.discrete.full_factorial.
FullFactorialGenerator
(max_cardinality: int = 100, check_cardinality: bool = True)[source]¶ Bases:
ax.models.discrete_base.DiscreteModel
Generator for full factorial designs.
Generates arms for all possible combinations of parameter values, each with weight 1.
The value of n supplied to gen will be ignored, as the number of arms generated is determined by the list of parameter values. To suppress this warning, use n = 1.

gen
(n: int, parameter_values: List[List[Optional[Union[str, bool, float, int]]]], objective_weights: Optional[numpy.ndarray], outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, Optional[Union[str, bool, float, int]]]] = None, pending_observations: Optional[List[List[List[Optional[Union[str, bool, float, int]]]]]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Tuple[List[List[Optional[Union[str, bool, float, int]]]], List[float], Dict[str, Any]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
parameter_values – A list of possible values for each parameter.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m lists of parameterizations (each parameterization is a list of parameter values of length d), each of length k_i, for each outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
 Returns
2element tuple containing
List of n generated points, where each point is represented by a list of parameter values.
List of weights for each of the n points.

ax.models.discrete.thompson module¶

class
ax.models.discrete.thompson.
ThompsonSampler
(num_samples: int = 10000, min_weight: Optional[float] = None, uniform_weights: bool = False)[source]¶ Bases:
ax.models.discrete_base.DiscreteModel
Generator for Thompson sampling.
The generator performs Thompson sampling on the data passed in via fit. Arms are given weight proportional to the probability that they are winners, according to Monte Carlo simulations.

fit
(Xs: List[List[List[Optional[Union[str, bool, float, int]]]]], Ys: List[List[float]], Yvars: List[List[float]], parameter_values: List[List[Optional[Union[str, bool, float, int]]]], outcome_names: List[str]) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m lists X of parameterizations (each parameterization is a list of parameter values of length d), each of length k_i, for each outcome.
Ys – The corresponding list of m lists Y, each of length k_i, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
parameter_values – A list of possible values for each parameter.
outcome_names – A list of m outcome names.

gen
(n: int, parameter_values: List[List[Optional[Union[str, bool, float, int]]]], objective_weights: Optional[numpy.ndarray], outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, Optional[Union[str, bool, float, int]]]] = None, pending_observations: Optional[List[List[List[Optional[Union[str, bool, float, int]]]]]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Tuple[List[List[Optional[Union[str, bool, float, int]]]], List[float], Dict[str, Any]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
parameter_values – A list of possible values for each parameter.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m lists of parameterizations (each parameterization is a list of parameter values of length d), each of length k_i, for each outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
 Returns
2element tuple containing
List of n generated points, where each point is represented by a list of parameter values.
List of weights for each of the n points.

predict
(X: List[List[Optional[Union[str, bool, float, int]]]]) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Predict
 Parameters
X – List of the j parameterizations at which to make predictions.
 Returns
2element tuple containing
(j x m) array of outcome predictions at X.
(j x m x m) array of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

NumPy Models¶
ax.models.numpy.randomforest module¶

class
ax.models.numpy.randomforest.
RandomForest
(max_features: Optional[str] = 'sqrt', num_trees: int = 500)[source]¶ Bases:
ax.models.numpy_base.NumpyModel
A Random Forest model.
Uses a parametric bootstrap to handle uncertainty in Y.
Can be used to fit data, make predictions, and do cross validation; however gen is not implemented and so this model cannot generate new points.
 Parameters
max_features – Maximum number of features at each split. With onehot encoding, this should be set to None. Defaults to “sqrt”, which is Breiman’s version of Random Forest.
num_trees – Number of trees.

cross_validate
(Xs_train: List[numpy.ndarray], Ys_train: List[numpy.ndarray], Yvars_train: List[numpy.ndarray], X_test: numpy.ndarray) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m (k_i x d) feature matrices X. Number of rows k_i can vary from i=1,…,m.
Ys_train – The corresponding list of m (k_i x 1) outcome arrays Y, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – (j x d) array of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) array of outcome predictions at X.
(j x m x m) array of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

fit
(Xs: List[numpy.ndarray], Ys: List[numpy.ndarray], Yvars: List[numpy.ndarray], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature matrices X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome arrays Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

predict
(X: numpy.ndarray) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Predict
 Parameters
X – (j x d) array of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) array of outcome predictions at X.
(j x m x m) array of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].
Random Models¶
ax.models.random.alebo_initializer module¶

class
ax.models.random.alebo_initializer.
ALEBOInitializer
(B: numpy.ndarray, nsamp: int = 10000, init_bound: int = 16, **kwargs: Any)[source]¶ Bases:
ax.models.random.uniform.UniformGenerator
Sample in a lowdimensional linear embedding, to initialize ALEBO.
Generates points on a linear subspace of [1, 1]^D by generating points in [b, b]^D, projecting them down with a matrix B, and then projecting them back up with the pseudoinverse of B. Thus points thus all lie in a linear subspace defined by B. Points whose upprojection falls outside of [1, 1]^D are thrown out, via rejection sampling.
To generate n points, we start with nsamp points in [b, b]^D, which are mapped down to the embedding and back up as described above. If >=n points fall within [1, 1]^D after being mapped up, then the first n are returned. If there are less than n points in [1, 1]^D, then b is constricted (halved) and the process is repeated until there are at least n points in [1, 1]^D. There exists a b small enough that all points will project to [1, 1]^D, so this is guaranteed to terminate, typically after few rounds.
 Parameters
B – A (dxD) projection down.
nsamp – Number of samples to use for rejection sampling.
init_bound – b for the initial sampling space described above.
kwargs – kwargs for UniformGenerator

gen
(n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X. Defined on [0, 1]^d.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
model_gen_options – A config dictionary that is passed along to the model.
rounding_func – A function that rounds an optimization result appropriately (e.g., according to roundtrip transformations).
 Returns
2element tuple containing
(n x d) array of generated points.
Uniform weights, an narray of ones for each point.
ax.models.random.base module¶

class
ax.models.random.base.
RandomModel
(deduplicate: bool = True, seed: Optional[int] = None, generated_points: Optional[numpy.ndarray] = None, fallback_to_sample_polytope: bool = False)[source]¶ Bases:
ax.models.base.Model
This class specifies the basic skeleton for a random model.
As random generators do not make use of models, they do not implement the fit or predict methods.
These models do not need data, or optimization configs.
To satisfy search space parameter constraints, these models can use rejection sampling. To enable rejection sampling for a subclass, only only _gen_samples needs to be implemented, or alternatively, _gen_unconstrained/gen can be directly implemented.

deduplicate
¶ If True (defaults to True), a single instantiation of the model will not return the same point twice. This flag is used in rejection sampling.

scramble
¶ If True, permutes the parameter values among the elements of the Sobol sequence. Default is True.

seed
¶ An optional seed value for scrambling.

gen
(n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X. Defined on [0, 1]^d.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
model_gen_options – A config dictionary that is passed along to the model.
rounding_func – A function that rounds an optimization result appropriately (e.g., according to roundtrip transformations).
 Returns
2element tuple containing
(n x d) array of generated points.
Uniform weights, an narray of ones for each point.

ax.models.random.rembo_initializer module¶

class
ax.models.random.rembo_initializer.
REMBOInitializer
(A: numpy.ndarray, bounds_d: List[Tuple[float, float]], **kwargs: Any)[source]¶ Bases:
ax.models.random.uniform.UniformGenerator
Sample in a lowdimensional linear embedding.
Generates points in [1, 1]^D by generating points in a ddimensional embedding, with box bounds as specified. When points are projected up, if they fall outside [1, 1]^D they are clamped to those bounds.
 Parameters
A – A (Dxd) linear embedding
bounds_d – Box bounds in the lowd space
kwargs – kwargs for UniformGenerator

gen
(n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X. Defined on [0, 1]^d.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
model_gen_options – A config dictionary that is passed along to the model.
rounding_func – A function that rounds an optimization result appropriately (e.g., according to roundtrip transformations).
 Returns
2element tuple containing
(n x d) array of generated points.
Uniform weights, an narray of ones for each point.
ax.models.random.sobol module¶

class
ax.models.random.sobol.
SobolGenerator
(seed: Optional[int] = None, deduplicate: bool = False, init_position: int = 0, scramble: bool = True, generated_points: Optional[numpy.ndarray] = None, fallback_to_sample_polytope: bool = False)[source]¶ Bases:
ax.models.random.base.RandomModel
This class specifies the generation algorithm for a Sobol generator.
As Sobol does not make use of a model, it does not implement the fit or predict methods.

deduplicate
¶ If true, a single instantiation of the generator will not return the same point twice.

init_position
¶ The initial state of the Sobol generator. Starts at 0 by default.

scramble
¶ If True, permutes the parameter values among the elements of the Sobol sequence. Default is True.

seed
¶ An optional seed value for scrambling.

property
engine
¶ Return a singleton SobolEngine.

gen
(n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None) → Tuple[numpy.ndarray, numpy.ndarray][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
rounding_func – A function that rounds an optimization result appropriately (e.g., according to roundtrip transformations) but unused here.
 Returns
2element tuple containing
(n x d) array of generated points.
Uniform weights, an narray of ones for each point.

ax.models.random.uniform module¶

class
ax.models.random.uniform.
UniformGenerator
(deduplicate: bool = False, seed: Optional[int] = None)[source]¶ Bases:
ax.models.random.base.RandomModel
This class specifies a uniform random generation algorithm.
As a uniform generator does not make use of a model, it does not implement the fit or predict methods.

seed
¶ An optional seed value for the underlying PRNG.

Torch Models¶
ax.models.torch.alebo module¶

class
ax.models.torch.alebo.
ALEBO
(B: torch.Tensor, laplace_nsamp: int = 25, fit_restarts: int = 10)[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
Does Bayesian optimization in a linear subspace with ALEBO.
The (d x D) projection down matrix B must be provided, and must be that used for the initialization.
Function evaluations happen in the highD space. We only evaluate points such that x = pinverse(B) @ B @ x (that is, points inside the subspace). Under that constraint, the projection is invertible.
 Parameters
B – (d x D) projection matrix (projects down).
laplace_nsamp – Number of samples for posterior sampling of kernel hyperparameters.
fit_restarts – Number of random restarts for MAP estimation.

Xs
: List[Tensor]¶

Ys
: List[Tensor]¶

Yvars
: List[Tensor]¶

best_point
(bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ Identify the current best point, satisfying the constraints in the same format as to gen.
Return None if no such point can be identified.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
dtensor of the best point.

cross_validate
(Xs_train: List[torch.Tensor], Ys_train: List[torch.Tensor], Yvars_train: List[torch.Tensor], X_test: torch.Tensor, **kwargs: Any) → Tuple[torch.Tensor, torch.Tensor][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys_train – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – (j x d) tensor of the j points at which to make predictions.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

fit
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

gen
(n: int, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], List[Optional[Dict[str, Any]]]][source]¶ Generate candidates.
Candidates are generated in the linear embedding with the polytope constraints described in the paper.
model_gen_options can contain ‘raw_samples’ (number of samples used for initializing the acquisition function optimization) and ‘num_restarts’ (number of restarts for acquisition function optimization).

get_and_fit_model
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], state_dicts: Optional[List[MutableMapping[str, torch.Tensor]]] = None) → botorch.models.gpytorch.GPyTorchModel[source]¶ Get a fitted ALEBO model for each outcome.
 Parameters
Xs – X for each outcome, already projected down.
Ys – Y for each outcome.
Yvars – Noise variance of Y for each outcome.
state_dicts – State dicts to initialize model fitting.
Returns: Fitted ALEBO model.

model
: Optional[Model]¶

predict
(X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predict
 Parameters
X – (j x d) tensor of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

update
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, **kwargs: Any) → None[source]¶ Update the model.
Updating the model requires both existing and additional data. The data passed into this method will become the new training data.
 Parameters
Xs – Existing + additional data for the model, in the same format as for fit.
Ys – Existing + additional data for the model, in the same format as for fit.
Yvars – Existing + additional data for the model, in the same format as for fit.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

class
ax.models.torch.alebo.
ALEBOGP
(B: torch.Tensor, train_X: torch.Tensor, train_Y: torch.Tensor, train_Yvar: torch.Tensor)[source]¶ Bases:
botorch.models.gp_regression.FixedNoiseGP
The GP for ALEBO.
Uses the Mahalanobis kernel defined in ALEBOKernel, along with a ScaleKernel to add a kernel variance and a fitted constant mean.
In nonbatch mode, there is a single kernel that produces MVN predictions as usual for a GP. With b batches, each batch has its own set of kernel hyperparameters and each batch represents a sample from the hyperparameter posterior distribution. When making a prediction (with __call__), these samples are integrated over using moment matching. So, the predictions are an MVN as usual with the same shape as in nonbatch mode.
 Parameters
B – (d x D) Projection matrix.
train_X – (n x d) X training data.
train_Y – (n x 1) Y training data.
train_Yvar – (n x 1) Noise variances of each training Y.

posterior
(X: torch.Tensor, output_indices: Optional[List[int]] = None, observation_noise: Union[bool, torch.Tensor] = False, **kwargs: Any) → botorch.posteriors.gpytorch.GPyTorchPosterior[source]¶ Computes the posterior over model outputs at the provided points.
 Parameters
X – A (batch_shape) x q x ddim Tensor, where d is the dimension of the feature space and q is the number of points considered jointly.
output_indices – A list of indices, corresponding to the outputs over which to compute the posterior (if the model is multioutput). Can be used to speed up computation if only a subset of the model’s outputs are required for optimization. If omitted, computes the posterior over all model outputs.
observation_noise – If True, add the observation noise from the likelihood to the posterior. If a Tensor, use it directly as the observation noise (must be of shape (batch_shape) x q x m).
 Returns
A GPyTorchPosterior object, representing batch_shape joint distributions over q points and the outputs selected by output_indices each. Includes observation noise if specified.

class
ax.models.torch.alebo.
ALEBOKernel
(B: torch.Tensor, batch_shape: torch.Size)[source]¶ Bases:
gpytorch.kernels.kernel.Kernel
The kernel for ALEBO.
Suppose there exists an ARD RBF GP on an (unknown) linear embedding with projection matrix A. We make function evaluations in a different linear embedding with projection matrix B (known). This is the appropriate kernel for fitting those data.
This kernel computes a Mahalanobis distance, and the (d x d) PD distance matrix Gamma is a parameter that must be fit. This is done by fitting its upper Cholesky decomposition, U.
 Parameters
B – (d x D) Projection matrix.
batch_shape – Batch shape as usual for gpytorch kernels.

ax.models.torch.alebo.
alebo_acqf_optimizer
(acq_function: botorch.acquisition.acquisition.AcquisitionFunction, bounds: torch.Tensor, n: int, inequality_constraints: Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], fixed_features: Optional[Dict[int, float]], rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]], raw_samples: int, num_restarts: int, B: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Optimize the acquisition function for ALEBO.
We are optimizing over a polytope within the subspace, and so begin each random restart of the acquisition function optimization with points that lie within that polytope.

ax.models.torch.alebo.
ei_or_nei
(model: Union[ax.models.torch.alebo.ALEBOGP, botorch.models.model_list_gp_regression.ModelListGP], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]], X_observed: torch.Tensor, X_pending: Optional[torch.Tensor], q: int, noiseless: bool) → botorch.acquisition.acquisition.AcquisitionFunction[source]¶ Use analytic EI if appropriate, otherwise Monte Carlo NEI.
Analytic EI can be used if: Single outcome, no constraints, no pending points, not batch, and no noise.
 Parameters
model – GP.
objective_weights – Weights on each outcome for the objective.
outcome_constraints – Outcome constraints.
X_observed – Observed points for NEI.
X_pending – Pending points.
q – Batch size.
noiseless – True if evaluations are noiseless.
Returns: An AcquisitionFunction, either analytic EI or MC NEI.

ax.models.torch.alebo.
extract_map_statedict
(m_b: Union[ax.models.torch.alebo.ALEBOGP, botorch.models.model_list_gp_regression.ModelListGP], num_outputs: int) → List[MutableMapping[str, torch.Tensor]][source]¶ Extract MAP statedict from the batchmode ALEBO GP.
The batch GP can be either a single ALEBO GP or a ModelListGP of ALEBO GPs.
 Parameters
m_b – Batchmode GP.
num_outputs – Number of outputs being modeled.

ax.models.torch.alebo.
get_batch_model
(B: torch.Tensor, train_X: torch.Tensor, train_Y: torch.Tensor, train_Yvar: torch.Tensor, Uvec_batch: torch.Tensor, mean_constant_batch: torch.Tensor, output_scale_batch: torch.Tensor) → ax.models.torch.alebo.ALEBOGP[source]¶ Construct a batchmode ALEBO GP using batch tensors of hyperparameters.
 Parameters
B – Projection matrix.
train_X – X training data.
train_Y – Y training data.
train_Yvar – Noise variances of each training Y.
Uvec_batch – Batch tensor of Uvec hyperparameters.
mean_constant_batch – Batch tensor of mean constant hyperparameter.
output_scale_batch – Batch tensor of output scale hyperparameter.
Returns: Batchmode ALEBO GP.

ax.models.torch.alebo.
get_fitted_model
(B: torch.Tensor, train_X: torch.Tensor, train_Y: torch.Tensor, train_Yvar: torch.Tensor, restarts: int, nsamp: int, init_state_dict: Optional[Dict[str, torch.Tensor]]) → ax.models.torch.alebo.ALEBOGP[source]¶ Get a fitted ALEBO GP.
We do random restart optimization to get a MAP model, then use the Laplace approximation to draw posterior samples of kernel hyperparameters, and finally construct a batchmode model where each batch is one of those sampled sets of kernel hyperparameters.
 Parameters
B – Projection matrix.
train_X – X training data.
train_Y – Y training data.
train_Yvar – Noise variances of each training Y.
restarts – Number of restarts for MAP estimation.
nsamp – Number of samples to draw from kernel hyperparameter posterior.
init_state_dict – Optionally begin MAP estimation with this state dict.
Returns: Batchmode (nsamp batches) fitted ALEBO GP.

ax.models.torch.alebo.
get_map_model
(B: torch.Tensor, train_X: torch.Tensor, train_Y: torch.Tensor, train_Yvar: torch.Tensor, restarts: int, init_state_dict: Optional[Dict[str, torch.Tensor]]) → gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood[source]¶ Do randomrestart optimization for MAP fitting of an ALEBO GP model.
 Parameters
B – Projection matrix.
train_X – X training data.
train_Y – Y training data.
train_Yvar – Noise variances of each training Y.
restarts – Number of restarts for MAP estimation.
init_state_dict – Optionally begin MAP estimation with this state dict.
Returns: nonbatch ALEBO GP with MAP kernel hyperparameters.

ax.models.torch.alebo.
laplace_sample_U
(mll: gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood, nsamp: int) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor][source]¶ Draw posterior samples of kernel hyperparameters using Laplace approximation.
Only the Mahalanobis distance matrix is sampled.
The diagonal of the Hessian is estimated using finite differences of the autograd gradients. The Laplace approximation is then N(p_map, inv(H)). We construct a set of nsamp kernel hyperparameters by drawing nsamp1 values from this distribution, and prepending as the first sample the MAP parameters.
 Parameters
mll – MLL object of MAP ALEBO GP.
nsamp – Number of samples to return.
 Returns: Batch tensors of the kernel hyperparameters Uvec, mean constant,
and output scale.
ax.models.torch.botorch module¶

class
ax.models.torch.botorch.
BotorchModel
(model_constructor: Callable[[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[int], List[int], List[str], Optional[Dict[str, torch.Tensor]], Any], botorch.models.model.Model] = <function get_and_fit_model>, model_predictor: Callable[[botorch.models.model.Model, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]] = <function predict_from_model>, acqf_constructor: Callable[[botorch.models.model.Model, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], Any], botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEI>, acqf_optimizer: Callable[[botorch.acquisition.acquisition.AcquisitionFunction, torch.Tensor, int, Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[Dict[int, float]], Optional[Callable[[torch.Tensor], torch.Tensor]], Any], Tuple[torch.Tensor, torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: Callable[[ax.models.torch_base.TorchModel, List[Tuple[float, float]], torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Dict[int, float]], Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]], Optional[Dict[int, float]]], Optional[torch.Tensor]] = <function recommend_best_observed_point>, refit_on_cv: bool = False, refit_on_update: bool = True, warm_start_refitting: bool = True, use_input_warping: bool = False, use_loocv_pseudo_likelihood: bool = False, **kwargs: Any)[source]¶ Bases:
ax.models.torch_base.TorchModel
Customizable botorch model.
By default, this uses a noisy Expected Improvement acquisition function on top of a model made up of separate GPs, one for each outcome. This behavior can be modified by providing custom implementations of the following components:
a model_constructor that instantiates and fits a model on data
a model_predictor that predicts outcomes using the fitted model
a acqf_constructor that creates an acquisition function from a fitted model
a acqf_optimizer that optimizes the acquisition function
 a best_point_recommender that recommends a current “best” point (i.e.,
what the model recommends if the learning process ended now)
 Parameters
model_constructor – A callable that instantiates and fits a model on data, with signature as described below.
model_predictor – A callable that predicts using the fitted model, with signature as described below.
acqf_constructor – A callable that creates an acquisition function from a fitted model, with signature as described below.
acqf_optimizer – A callable that optimizes the acquisition function, with signature as described below.
best_point_recommender – A callable that recommends the best point, with signature as described below.
refit_on_cv – If True, refit the model for each fold when performing crossvalidation.
refit_on_update – If True, refit the model after updating the training data using the update method.
warm_start_refitting – If True, start model refitting from previous model parameters in order to speed up the fitting process.
Call signatures:
model_constructor( Xs, Ys, Yvars, task_features, fidelity_features, metric_names, state_dict, **kwargs, ) > model
Here Xs, Ys, Yvars are lists of tensors (one element per outcome), task_features identifies columns of Xs that should be modeled as a task, fidelity_features is a list of ints that specify the positions of fidelity parameters in ‘Xs’, metric_names provides the names of each Y in Ys, state_dict is a pytorch module state dict, and model is a BoTorch Model. Optional kwargs are being passed through from the BotorchModel constructor. This callable is assumed to return a fitted BoTorch model that has the same dtype and lives on the same device as the input tensors.
model_predictor(model, X) > [mean, cov]
Here model is a fitted botorch model, X is a tensor of candidate points, and mean and cov are the posterior mean and covariance, respectively.
acqf_constructor( model, objective_weights, outcome_constraints, X_observed, X_pending, **kwargs, ) > acq_function
Here model is a botorch Model, objective_weights is a tensor of weights for the model outputs, outcome_constraints is a tuple of tensors describing the (linear) outcome constraints, X_observed are previously observed points, and X_pending are points whose evaluation is pending. acq_function is a BoTorch acquisition function crafted from these inputs. For additional details on the arguments, see get_NEI.
acqf_optimizer( acq_function, bounds, n, inequality_constraints, equality_constraints, fixed_features, rounding_func, **kwargs, ) > candidates
Here acq_function is a BoTorch AcquisitionFunction, bounds is a tensor containing bounds on the parameters, n is the number of candidates to be generated, inequality_constraints are inequality constraints on parameter values, fixed_features specifies features that should be fixed during generation, and rounding_func is a callback that rounds an optimization result appropriately. candidates is a tensor of generated candidates. For additional details on the arguments, see scipy_optimizer.
best_point_recommender( model, bounds, objective_weights, outcome_constraints, linear_constraints, fixed_features, model_gen_options, target_fidelities, ) > candidates
Here model is a TorchModel, bounds is a list of tuples containing bounds on the parameters, objective_weights is a tensor of weights for the model outputs, outcome_constraints is a tuple of tensors describing the (linear) outcome constraints, linear_constraints is a tuple of tensors describing constraints on the design, fixed_features specifies features that should be fixed during generation, model_gen_options is a config dictionary that can contain modelspecific options, and target_fidelities is a map from fidelity feature column indices to their respective target fidelities, used for multifidelity optimization problems. % TODO: refer to an example.

Xs
: List[torch.Tensor]¶

Ys
: List[torch.Tensor]¶

Yvars
: List[torch.Tensor]¶

best_point
(bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ Identify the current best point, satisfying the constraints in the same format as to gen.
Return None if no such point can be identified.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
dtensor of the best point.

cross_validate
(Xs_train: List[torch.Tensor], Ys_train: List[torch.Tensor], Yvars_train: List[torch.Tensor], X_test: torch.Tensor, **kwargs: Any) → Tuple[torch.Tensor, torch.Tensor][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys_train – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – (j x d) tensor of the j points at which to make predictions.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

fit
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

gen
(n: int, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], Optional[List[Optional[Dict[str, Any]]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
4element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
Generation metadata
 Dictionary of modelspecific metadata for the given
generation candidates

model
: Optional[botorch.models.model.Model]¶

predict
(X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predict
 Parameters
X – (j x d) tensor of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

update
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, **kwargs: Any) → None[source]¶ Update the model.
Updating the model requires both existing and additional data. The data passed into this method will become the new training data.
 Parameters
Xs – Existing + additional data for the model, in the same format as for fit.
Ys – Existing + additional data for the model, in the same format as for fit.
Yvars – Existing + additional data for the model, in the same format as for fit.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.
ax.models.torch.botorch_defaults module¶

ax.models.torch.botorch_defaults.
get_NEI
(model: botorch.models.model.Model, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None, X_pending: Optional[torch.Tensor] = None, **kwargs: Any) → botorch.acquisition.acquisition.AcquisitionFunction[source]¶ Instantiates a qNoisyExpectedImprovement acquisition function.

ax.models.torch.botorch_defaults.
get_and_fit_model
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, torch.Tensor]] = None, refit_model: bool = True, use_input_warping: bool = False, use_loocv_pseudo_likelihood: bool = False, **kwargs: Any) → botorch.models.gpytorch.GPyTorchModel[source]¶ Instantiates and fits a botorch GPyTorchModel using the given data. N.B. Currently, the logic for choosing ModelListGP vs other models is handled using ifelse statements in lines 96137. In the future, this logic should be taken care of by modular botorch.
 Parameters
Xs – List of X data, one tensor per outcome.
Ys – List of Y data, one tensor per outcome.
Yvars – List of observed variance of Ys.
task_features – List of columns of X that are tasks.
fidelity_features – List of columns of X that are fidelity parameters.
metric_names – Names of each outcome Y in Ys.
state_dict – If provided, will set model parameters to this state dictionary. Otherwise, will fit the model.
refit_model – Flag for refitting model.
 Returns
A fitted GPyTorchModel.

ax.models.torch.botorch_defaults.
get_warping_transform
(d: int, batch_shape: Optional[torch.Size] = None, task_feature: Optional[int] = None) → botorch.models.transforms.input.Warp[source]¶ Construct input warping transform.
 Parameters
d – The dimension of the input, including task features
batch_shape – The batch_shape of the model
task_feature – The index of the task feature
 Returns
The input warping transform.

ax.models.torch.botorch_defaults.
recommend_best_observed_point
(model: ax.models.torch_base.TorchModel, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ A wrapper around ax.models.model_utils.best_observed_point for TorchModel that recommends a best point from previously observed points using either a “max_utility” or “feasible_threshold” strategy.
 Parameters
model – A TorchModel.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
A darray of the best point, or None if no feasible point was observed.

ax.models.torch.botorch_defaults.
recommend_best_out_of_sample_point
(model: ax.models.torch_base.TorchModel, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ Identify the current best point by optimizing the posterior mean of the model. This is “outofsample” because it considers unobserved designs as well.
Return None if no such point can be identified.
 Parameters
model – A TorchModel.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
A darray of the best point, or None if no feasible point exists.

ax.models.torch.botorch_defaults.
scipy_optimizer
(acq_function: botorch.acquisition.acquisition.AcquisitionFunction, bounds: torch.Tensor, n: int, inequality_constraints: Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]] = None, equality_constraints: Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, **kwargs: Any) → Tuple[torch.Tensor, torch.Tensor][source]¶ Optimizer using scipy’s minimize module on a numpyadpator.
 Parameters
acq_function – A botorch AcquisitionFunction.
bounds – A 2 x ddim tensor, where bounds[0] (bounds[1]) are the lower (upper) bounds of the feasible hyperrectangle.
n – The number of candidates to generate.
constraints (equality) – A list of tuples (indices, coefficients, rhs), with each tuple encoding an inequality constraint of the form sum_i (X[indices[i]] * coefficients[i]) >= rhs
constraints – A list of tuples (indices, coefficients, rhs), with each tuple encoding an equality constraint of the form sum_i (X[indices[i]] * coefficients[i]) == rhs
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
 Returns
2element tuple containing
A n x ddim tensor of generated candidates.
In the case of joint optimization, a scalar tensor containing the joint acquisition value of the n points. In the case of sequential optimization, a ndim tensor of conditional acquisition values, where ith element is the expected acquisition value conditional on having observed candidates 0,1,…,i1.
ax.models.torch.botorch_kg module¶

class
ax.models.torch.botorch_kg.
KnowledgeGradient
(cost_intercept: float = 1.0, linear_truncated: bool = True, use_input_warping: bool = False, **kwargs: Any)[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
The Knowledge Gradient with one shot optimization.
 Parameters
cost_intercept – The cost intercept for the affine cost of the form cost_intercept + n, where n is the number of generated points. Only used for multifidelity optimzation (i.e., if fidelity_features are present).
linear_truncated – If False, use an alternate downsampling + exponential decay Kernel instead of the default LinearTruncatedFidelityKernel (only relevant for multifidelity optimization).
kwargs – Modelspecific kwargs.

Xs
: List[torch.Tensor]¶

Ys
: List[torch.Tensor]¶

Yvars
: List[torch.Tensor]¶

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

gen
(n: int, bounds: List, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], Optional[List[Optional[Dict[str, Any]]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
3element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
 Dictionary of modelspecific metadata for the given
generation candidates.

model
: Optional[botorch.models.model.Model]¶
ax.models.torch.botorch_mes module¶

class
ax.models.torch.botorch_mes.
MaxValueEntropySearch
(cost_intercept: float = 1.0, linear_truncated: bool = True, use_input_warping: bool = False, **kwargs: Any)[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
Maxvalue entropy search.
 Parameters
cost_intercept – The cost intercept for the affine cost of the form cost_intercept + n, where n is the number of generated points. Only used for multifidelity optimzation (i.e., if fidelity_features are present).
linear_truncated – If False, use an alternate downsampling + exponential decay Kernel instead of the default LinearTruncatedFidelityKernel (only relevant for multifidelity optimization).
kwargs – Modelspecific kwargs.

Xs
: List[torch.Tensor]¶

Ys
: List[torch.Tensor]¶

Yvars
: List[torch.Tensor]¶

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

gen
(n: int, bounds: List, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], List[Optional[Dict[str, Any]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
4element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
Generation metadata
 Dictionary of modelspecific metadata for the given
generation candidates

model
: Optional[botorch.models.model.Model]¶
ax.models.torch.botorch_moo module¶

class
ax.models.torch.botorch_moo.
MultiObjectiveBotorchModel
(model_constructor: Callable[[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[int], List[int], List[str], Optional[Dict[str, torch.Tensor]], Any], botorch.models.model.Model] = <function get_and_fit_model>, model_predictor: Callable[[botorch.models.model.Model, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]] = <function predict_from_model>, acqf_constructor: Callable[[botorch.models.model.Model, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], Any], botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEHVI>, acqf_optimizer: Callable[[botorch.acquisition.acquisition.AcquisitionFunction, torch.Tensor, int, Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[Dict[int, float]], Optional[Callable[[torch.Tensor], torch.Tensor]], Any], Tuple[torch.Tensor, torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: Callable[[ax.models.torch_base.TorchModel, List[Tuple[float, float]], torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Dict[int, float]], Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]], Optional[Dict[int, float]]], Optional[torch.Tensor]] = <function recommend_best_observed_point>, frontier_evaluator: Callable[[ax.models.torch_base.TorchModel, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = <function pareto_frontier_evaluator>, refit_on_cv: bool = False, refit_on_update: bool = True, warm_start_refitting: bool = False, use_input_warping: bool = False, use_loocv_pseudo_likelihood: bool = False, **kwargs: Any)[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
Customizable multiobjective model.
By default, this uses an Expected Hypervolume Improvment function to find the pareto frontier of a function with multiple outcomes. This behavior can be modified by providing custom implementations of the following components:
a model_constructor that instantiates and fits a model on data
a model_predictor that predicts outcomes using the fitted model
a acqf_constructor that creates an acquisition function from a fitted model
a acqf_optimizer that optimizes the acquisition function
 Parameters
model_constructor – A callable that instantiates and fits a model on data, with signature as described below.
model_predictor – A callable that predicts using the fitted model, with signature as described below.
acqf_constructor – A callable that creates an acquisition function from a fitted model, with signature as described below.
acqf_optimizer – A callable that optimizes an acquisition function, with signature as described below.
Call signatures:
model_constructor( Xs, Ys, Yvars, task_features, fidelity_features, metric_names, state_dict, **kwargs, ) > model
Here Xs, Ys, Yvars are lists of tensors (one element per outcome), task_features identifies columns of Xs that should be modeled as a task, fidelity_features is a list of ints that specify the positions of fidelity parameters in ‘Xs’, metric_names provides the names of each Y in Ys, state_dict is a pytorch module state dict, and model is a BoTorch Model. Optional kwargs are being passed through from the BotorchModel constructor. This callable is assumed to return a fitted BoTorch model that has the same dtype and lives on the same device as the input tensors.
model_predictor(model, X) > [mean, cov]
Here model is a fitted botorch model, X is a tensor of candidate points, and mean and cov are the posterior mean and covariance, respectively.
acqf_constructor( model, objective_weights, outcome_constraints, X_observed, X_pending, **kwargs, ) > acq_function
Here model is a botorch Model, objective_weights is a tensor of weights for the model outputs, outcome_constraints is a tuple of tensors describing the (linear) outcome constraints, X_observed are previously observed points, and X_pending are points whose evaluation is pending. acq_function is a BoTorch acquisition function crafted from these inputs. For additional details on the arguments, see get_NEI.
acqf_optimizer( acq_function, bounds, n, inequality_constraints, fixed_features, rounding_func, **kwargs, ) > candidates
Here acq_function is a BoTorch AcquisitionFunction, bounds is a tensor containing bounds on the parameters, n is the number of candidates to be generated, inequality_constraints are inequality constraints on parameter values, fixed_features specifies features that should be fixed during generation, and rounding_func is a callback that rounds an optimization result appropriately. candidates is a tensor of generated candidates. For additional details on the arguments, see scipy_optimizer.
frontier_evaluator( model, objective_weights, objective_thresholds, X, Y, Yvar, outcome_constraints, )
Here model is a botorch Model, objective_thresholds is used in hypervolume evaluations, objective_weights is a tensor of weights applied to the objectives (sign represents direction), X, Y, Yvar are tensors, outcome_constraints is a tuple of tensors describing the (linear) outcome constraints.

Xs
: List[torch.Tensor]¶

Ys
: List[torch.Tensor]¶

Yvars
: List[torch.Tensor]¶

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

gen
(n: int, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, objective_thresholds: Optional[torch.Tensor] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], Optional[List[Optional[Dict[str, Any]]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
4element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
Generation metadata
 Dictionary of modelspecific metadata for the given
generation candidates
ax.models.torch.botorch_moo_defaults module¶
References
 Daulton2020qehvi
S. Daulton, M. Balandat, and E. Bakshy. Differentiable Expected Hypervolume Improvement for Parallel MultiObjective Bayesian Optimization. Advances in Neural Information Processing Systems 33, 2020.
 Daulton2021nehvi
S. Daulton, M. Balandat, and E. Bakshy. Parallel Bayesian Optimization of Multiple Noisy Objectives with Expected Hypervolume Improvement. Advances in Neural Information Processing Systems 34, 2021.

ax.models.torch.botorch_moo_defaults.
get_EHVI
(model: botorch.models.model.Model, objective_weights: torch.Tensor, objective_thresholds: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None, X_pending: Optional[torch.Tensor] = None, **kwargs: Any) → botorch.acquisition.acquisition.AcquisitionFunction[source]¶ Instantiates a qExpectedHyperVolumeImprovement acquisition function.
 Parameters
model – The underlying model which the acqusition function uses to estimate acquisition values of candidates.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
objective_thresholds – A tensor containing thresholds forming a reference point from which to calculate pareto frontier hypervolume. Points that do not dominate the objective_thresholds contribute nothing to hypervolume.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
X_observed – A tensor containing points observed for all objective outcomes and outcomes that appear in the outcome constraints (if there are any).
X_pending – A tensor containing points whose evaluation is pending (i.e. that have been submitted for evaluation) present for all objective outcomes and outcomes that appear in the outcome constraints (if there are any).
mc_samples – The number of MC samples to use (default: 512).
qmc – If True, use qMC instead of MC (default: True).
 Returns
The instantiated acquisition function.
 Return type
qExpectedHypervolumeImprovement

ax.models.torch.botorch_moo_defaults.
get_NEHVI
(model: botorch.models.model.Model, objective_weights: torch.Tensor, objective_thresholds: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None, X_pending: Optional[torch.Tensor] = None, **kwargs: Any) → botorch.acquisition.acquisition.AcquisitionFunction[source]¶ Instantiates a qNoisyExpectedHyperVolumeImprovement acquisition function.
 Parameters
model – The underlying model which the acqusition function uses to estimate acquisition values of candidates.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
X_observed – A tensor containing points observed for all objective outcomes and outcomes that appear in the outcome constraints (if there are any).
X_pending – A tensor containing points whose evaluation is pending (i.e. that have been submitted for evaluation) present for all objective outcomes and outcomes that appear in the outcome constraints (if there are any).
mc_samples – The number of MC samples to use (default: 512).
qmc – If True, use qMC instead of MC (default: True).
prune_baseline – If True, prune the baseline points for NEI (default: True).
chebyshev_scalarization – Use augmented Chebyshev scalarization.
 Returns
The instantiated acquisition function.
 Return type
qNoisyExpectedHyperVolumeImprovement

ax.models.torch.botorch_moo_defaults.
get_default_frontier_evaluator
() → Callable[[ax.models.torch_base.TorchModel, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]][source]¶

ax.models.torch.botorch_moo_defaults.
get_weighted_mc_objective_and_objective_thresholds
(objective_weights: torch.Tensor, objective_thresholds: torch.Tensor) → Tuple[botorch.acquisition.multi_objective.objective.WeightedMCMultiOutputObjective, torch.Tensor][source]¶ Construct weighted objective and apply the weights to objective thresholds.
 Parameters
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
objective_thresholds – A tensor containing thresholds forming a reference point from which to calculate pareto frontier hypervolume. Points that do not dominate the objective_thresholds contribute nothing to hypervolume.
 Returns
The objective
The objective thresholds
 Return type
A twoelement tuple with the objective and objective thresholds

ax.models.torch.botorch_moo_defaults.
infer_objective_thresholds
(model: botorch.models.model.Model, objective_weights: torch.Tensor, bounds: Optional[List[Tuple[float, float]]] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, subset_idcs: Optional[torch.Tensor] = None, Xs: Optional[List[torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None) → torch.Tensor[source]¶ Infer objective thresholds.
This method uses the modelestimated Pareto frontier over the insample points to infer absolute (not relativized) objective thresholds.
This uses a heuristic that sets the objective threshold to be a scaled nadir point, where the nadir point is scaled back based on the range of each objective across the current insample Pareto frontier.
See botorch.utils.multi_objective.hypervolume.infer_reference_point for details on the heuristic.
 Parameters
model – A fitted botorch Model.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights. These should not be subsetted.
bounds – A list of (lower, upper) tuples for each column of X.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. These should not be subsetted.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
subset_idcs – The indices of the outcomes that are modeled by the provided model. If subset_idcs not None, this method infers whether the model is subsetted.
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
X_observed – A n x ddim tensor of insample points to use for determining the current insample Pareto frontier.
 Returns
 A mdim tensor of objective thresholds, where the objective
threshold is nan if the outcome is not an objective.

ax.models.torch.botorch_moo_defaults.
pareto_frontier_evaluator
(model: ax.models.torch_base.TorchModel, objective_weights: torch.Tensor, objective_thresholds: Optional[torch.Tensor] = None, X: Optional[torch.Tensor] = None, Y: Optional[torch.Tensor] = None, Yvar: Optional[torch.Tensor] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor][source]¶ Return outcomes predicted to lie on a pareto frontier.
Given a model and points to evaluate, use the model to predict which points lie on the Pareto frontier.
 Parameters
model – Model used to predict outcomes.
objective_weights – A m tensor of values indicating the weight to put on different outcomes. For pareto frontiers only the sign matters.
objective_thresholds – A tensor containing thresholds forming a reference point from which to calculate pareto frontier hypervolume. Points that do not dominate the objective_thresholds contribute nothing to hypervolume.
X – A n x d tensor of features to evaluate.
Y – A n x m tensor of outcomes to use instead of predictions.
Yvar – A n x m x m tensor of input covariances (NaN if unobserved).
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
 Returns
3element tuple containing
 A j x m tensor of outcome on the pareto frontier. j is the number
of frontier points.
 A j x m x m tensor of predictive covariances.
cov[j, m1, m2] is Cov[[email protected], [email protected]].
A j tensor of the index of each frontier point in the input Y.

ax.models.torch.botorch_moo_defaults.
scipy_optimizer_list
(acq_function_list: List[botorch.acquisition.acquisition.AcquisitionFunction], bounds: torch.Tensor, inequality_constraints: Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, **kwargs: Any) → Tuple[torch.Tensor, torch.Tensor][source]¶ Sequential optimizer using scipy’s minimize module on a numpyadaptor.
The ith acquisition in the sequence uses the ith given acquisition_function.
 Parameters
acq_function_list – A list of botorch AcquisitionFunctions, optimized sequentially.
bounds – A 2 x ddim tensor, where bounds[0] (bounds[1]) are the lower (upper) bounds of the feasible hyperrectangle.
n – The number of candidates to generate.
constraints (inequality) – A list of tuples (indices, coefficients, rhs), with each tuple encoding an inequality constraint of the form sum_i (X[indices[i]] * coefficients[i]) >= rhs
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
 Returns
2element tuple containing
A n x ddim tensor of generated candidates.
A ndim tensor of conditional acquisition values, where ith element is the expected acquisition value conditional on having observed candidates 0,1,…,i1.
ax.models.torch.botorch_modular.acquisition module¶

class
ax.models.torch.botorch_modular.acquisition.
Acquisition
(surrogate: ax.models.torch.botorch_modular.surrogate.Surrogate, search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, botorch_acqf_class: Type[botorch.acquisition.acquisition.AcquisitionFunction], options: Optional[Dict[str, Any]] = None, pending_observations: Optional[List[torch.Tensor]] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, objective_thresholds: Optional[torch.Tensor] = None)[source]¶ Bases:
ax.utils.common.base.Base
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.
 Parameters
surrogate – Surrogate model, 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).
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
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.
pending_observations – A list of tensors, each of which contains points whose evaluation is pending (i.e. that have been submitted for evaluation) for a given outcome. A list of m (k_i x d) feature tensors X for m outcomes and k_i, pending observations for outcome i.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b. (Not used by single task models)
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.

acqf
: botorch.acquisition.acquisition.AcquisitionFunction¶

best_point
(search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Tuple[torch.Tensor, float][source]¶ Finds the best observed point and the corresponding observed outcome values.

property
botorch_acqf_class
¶ BoTorch
AcquisitionFunction
class underlying thisAcquisition
.

compute_model_dependencies
(surrogate: ax.models.torch.botorch_modular.surrogate.Surrogate, search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, pending_observations: Optional[List[torch.Tensor]] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Any]] = None) → Dict[str, Any][source]¶ Computes inputs to acquisition function class based on the given surrogate model.
NOTE: When subclassing Acquisition from a superclass where this method returns a nonempty dictionary of kwargs to AcquisitionFunction, call super().compute_model_dependencies and then update that dictionary of options with the options for the subclass you are creating (unless the superclass’ model dependencies should not be propagated to the subclass). See MultiFidelityKnowledgeGradient.compute_model_dependencies for an example.
 Parameters
surrogate – The surrogate object containing the BoTorch Model, with which this Acquisition is to be used.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
pending_observations – A list of tensors, each of which contains points whose evaluation is pending (i.e. that have been submitted for evaluation) for a given outcome. A list of m (k_i x d) feature tensors X for m outcomes and k_i, pending observations for outcome i.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b. (Not used by single task models)
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
options – The options kwarg dict, passed on initialization of the Acquisition object.
 Returns: A dictionary of surrogate modeldependent options, to be passed
as kwargs to BoTorch`AcquisitionFunction` constructor.

property
device
¶ Torch device type of the tensors in the training data used in the model, of which this
Acquisition
is a subcomponent.

property
dtype
¶ Torch data type of the tensors in the training data used in the model, of which this
Acquisition
is a subcomponent.

evaluate
(X: torch.Tensor) → torch.Tensor[source]¶ Evaluate the acquisition function on the candidate set X.
 Parameters
X – A batch_shape x q x ddim Tensor of tbatches with q ddim design points each.
 Returns
A batch_shape’dim Tensor of acquisition values at the given design points X, where batch_shape’ is the broadcasted batch shape of model and input X.

get_botorch_objective
(botorch_acqf_class: Type[botorch.acquisition.acquisition.AcquisitionFunction], model: botorch.models.model.Model, objective_weights: torch.Tensor, objective_thresholds: Optional[torch.Tensor] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None) → botorch.acquisition.objective.AcquisitionObjective[source]¶

property
objective_thresholds
¶ The objective thresholds for all outcomes.
For nonobjective outcomes, the objective thresholds are nans.

optimize
(n: int, search_space_digest: ax.core.search_space.SearchSpaceDigest, inequality_constraints: Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, optimizer_options: Optional[Dict[str, Any]] = None) → Tuple[torch.Tensor, torch.Tensor][source]¶ Generate a set of candidates via multistart optimization. Obtains candidates and their associated acquisition function values.
 Parameters
n – The number of candidates to generate.
search_space_digest – A
SearchSpaceDigest
object containing search space properties, e.g.bounds
for optimization.constraints (inequality) – A list of tuples (indices, coefficients, rhs), with each tuple encoding an inequality constraint of the form
sum_i (X[indices[i]] * coefficients[i]) >= rhs
.fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
rounding_func – A function that postprocesses an optimization result appropriately (i.e., according to roundtrip transformations).
optimizer_options – Options for the optimizer function, e.g.
sequential
orraw_samples
.
ax.models.torch.botorch_modular.default_options module¶

ax.models.torch.botorch_modular.default_options.
get_default_optimizer_options
(acqf_class: Type[botorch.acquisition.acquisition.AcquisitionFunction]) → Dict[str, Any][source]¶ Returns default options as a dict for a given acquisition function class.
NOTE: Logs a warning and returns
DEFAULT_OPTIMIZER_OPTIONS
if acquisition function is not registered inDEFAULT_OPTIMIZER_OPTIONS
.

ax.models.torch.botorch_modular.default_options.
mk_ehvi_default_optimizer_options
() → Dict[str, Any][source]¶ Makes a copy of dictionary for generic default optimizer options for EHVIbased acquisition function, used when optimizer options for a given acquisition function are not registered.
NOTE: Return of this function is safe to modify without affecting the default options returned subsequently.

ax.models.torch.botorch_modular.default_options.
mk_generic_default_optimizer_options
() → Dict[str, Any][source]¶ Makes a copy of dictionary for generic default optimizer options, used when optimizer options for a given acquisition function are not registered.
NOTE: Return of this function is safe to modify without affecting the default options returned subsequently.
ax.models.torch.botorch_modular.list_surrogate module¶

class
ax.models.torch.botorch_modular.list_surrogate.
ListSurrogate
(botorch_submodel_class_per_outcome: Optional[Dict[str, Type[botorch.models.model.Model]]] = None, botorch_submodel_class: Optional[Type[botorch.models.model.Model]] = None, submodel_options_per_outcome: Optional[Dict[str, Dict[str, Any]]] = None, submodel_options: Optional[Dict[str, Any]] = None, mll_class: Type[gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood] = <class 'gpytorch.mlls.sum_marginal_log_likelihood.SumMarginalLogLikelihood'>, mll_options: Optional[Dict[str, Any]] = None)[source]¶ Bases:
ax.models.torch.botorch_modular.surrogate.Surrogate
Special type of
Surrogate
that wraps a set of submodels intoModelListGP
under the hood for multioutcome or multitask models. Parameters
botorch_submodel_class_per_outcome – Mapping from metric name to BoTorch model class that should be used as surrogate model for that metric. Use instead of
botorch_submodel_class
.botorch_submodel_class – BoTorch
Model
class, shortcut for when all submodels of this surrogate’s underlyingModelListGP
are of the same type. Use instead ofbotorch_submodel_class_per_outcome
.submodel_options_per_outcome – Optional mapping from metric name to dictionary of kwargs for the submodel for that outcome.
submodel_options – Optional dictionary of kwargs, shared between all submodels. NOTE: kwargs for submodel are
submodel_options
(shared) +submodel_outions_per_outcome[submodel_outcome]
(individual).mll_class –
MarginalLogLikelihood
class to use for modelfitting.

botorch_submodel_class
: Optional[Type[botorch.models.model.Model]]¶

construct
(training_data: botorch.utils.containers.TrainingData, **kwargs: Any) → None[source]¶ Constructs the underlying BoTorch
Model
using the training data. Parameters
training_data – List of
TrainingData
for the submodels ofModelListGP
. Each training data is for one outcome, and the order of outcomes should match the order of metrics inmetric_names
argument.**kwargs – Keyword arguments, accepts: 
metric_names
(required): Names of metrics, in the same order as training data (so if training data is[tr_A, tr_B]
, the metrics are["A" and "B"]
). These are used to match training data with correct submodels ofModelListGP
, fidelity_features
: Indices of columns in X that represent fidelity, task_features
: Indices of columns in X that represent tasks.

property
device
¶

property
dtype
¶

kernel_class
: Optional[Type[gpytorch.kernels.kernel.Kernel]] = None¶

mll_class
: Type[gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood]¶

property
training_data_per_outcome
¶
ax.models.torch.botorch_modular.model module¶

class
ax.models.torch.botorch_modular.model.
BoTorchModel
(acquisition_class: Optional[Type[ax.models.torch.botorch_modular.acquisition.Acquisition]] = None, acquisition_options: Optional[Dict[str, Any]] = None, botorch_acqf_class: Optional[Type[botorch.acquisition.acquisition.AcquisitionFunction]] = None, surrogate: Optional[ax.models.torch.botorch_modular.surrogate.Surrogate] = None, surrogate_options: Optional[Dict[str, Any]] = None, refit_on_update: bool = True, refit_on_cv: bool = False, warm_start_refit: bool = True)[source]¶ Bases:
ax.models.torch_base.TorchModel
,ax.utils.common.base.Base
All classes in ‘botorch_modular’ directory are under construction, incomplete, and should be treated as alpha versions only.
Modular Model class for combining BoTorch subcomponents in Ax. Specified via Surrogate and Acquisition, which wrap BoTorch Model and AcquisitionFunction, respectively, for convenient use in Ax.
 Parameters
acquisition_class – Type of Acquisition to be used in this model, autoselected based on experiment and data if not specified.
acquisition_options – Optional dict of kwargs, passed to the constructor of BoTorch AcquisitionFunction.
botorch_acqf_class – Type of AcquisitionFunction to be used in this model, autoselected based on experiment and data if not specified.
surrogate – An instance of Surrogate to be used as part of this model; if not specified, type of Surrogate and underlying BoTorch Model will be autoselected based on experiment and data, with kwargs in surrogate_options applied.
surrogate_options – Optional dict of kwargs for Surrogate (used if no preinstantiated Surrogate via is passed via surrogate). Can include:  model_options: Dict of options to surrogate’s underlying BoTorch Model,  submodel_options or submodel_options_per_outcome: Options for submodels in ListSurrogate, see documentation for ListSurrogate.
refit_on_update – Whether to reoptimize model parameters during call to BoTorchModel.update. If false, training data for the model (used for inference) is still swapped for new training data, but model parameters are not reoptimized.
refit_on_cv – Whether to reoptimize model parameters during call to BoTorchmodel.cross_validate.
warm_start_refit – Whether to load parameters from either the provided state dict or the state dict of the current BoTorch Model during refitting. If False, model parameters will be reoptimized from scratch on refit. NOTE: This setting is ignored during update or cross_validate if the corresponding refit_on_… is False.

property
Xs
¶ A list of tensors, each of shape
batch_shape x n_i x d
, where n_i is the number of training inputs for the ith model.NOTE: This is an accessor for
self.surrogate.training_data.Xs
and returns it unchanged.

property
Ys
¶ A list of tensors, each of shape
batch_shape x n_i x 1
, where n_i is the number of training observations for the ith (singleoutput) model.NOTE: This is an accessor for
self.surrogate.training_data.Ys
and returns it unchanged.

property
Yvars
¶ An optional list of tensors, each of shape
batch_shape x n_i x 1
, wheren_i
is the number of training observations of the observation noise for the ith (singleoutput) model. If None, the observation noise level is unobserved.NOTE: This is an accessor for
self.surrogate.training_data.Yvars
and returns it unchanged.

acquisition_class
: Type[ax.models.torch.botorch_modular.acquisition.Acquisition]¶

best_point
(bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ Identify the current best point, satisfying the constraints in the same format as to gen.
Return None if no such point can be identified.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
dtensor of the best point.

property
botorch_acqf_class
¶ BoTorch
AcquisitionFunction
class, associated with this model. Raises an error if one is not yet set.

cross_validate
(Xs_train: List[torch.Tensor], Ys_train: List[torch.Tensor], Yvars_train: List[torch.Tensor], X_test: torch.Tensor, search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str]) → Tuple[torch.Tensor, torch.Tensor][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys_train – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – (j x d) tensor of the j points at which to make predictions.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

evaluate_acquisition_function
(X: torch.Tensor, search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, objective_thresholds: Optional[torch.Tensor] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, acq_options: Optional[Dict[str, Any]] = None) → torch.Tensor[source]¶ Evaluate the acquisition function on the candidate set X.
 Parameters
X – (j x d) tensor of the j points at which to evaluate the acquisition function.
search_space_digest – A dataclass used to compactly represent a search space.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
objective_thresholds – The mdim tensor of objective thresholds. There is one for each modeled metric.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be held fixed during the evaluation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
acq_options – Keyword arguments used to contruct the acquisition function.
 Returns
A singleelement tensor with the acquisition value for these points.

fit
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], target_fidelities: Optional[Dict[int, float]] = None, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, state_dict: Optional[Dict[str, torch.Tensor]] = None, refit: bool = True) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

gen
(n: int, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, objective_thresholds: Optional[torch.Tensor] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], Optional[List[Optional[Dict[str, Any]]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
4element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
Generation metadata
 Dictionary of modelspecific metadata for the given
generation candidates

predict
(X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predict
 Parameters
X – (j x d) tensor of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

property
surrogate
¶ Ax
Surrogate
object (wrapper for BoTorchModel
), associated with this model. Raises an error if one is not yet set.

update
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Update the model.
Updating the model requires both existing and additional data. The data passed into this method will become the new training data.
 Parameters
Xs – Existing + additional data for the model, in the same format as for fit.
Ys – Existing + additional data for the model, in the same format as for fit.
Yvars – Existing + additional data for the model, in the same format as for fit.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.
ax.models.torch.botorch_modular.multi_fidelity module¶

class
ax.models.torch.botorch_modular.multi_fidelity.
MultiFidelityAcquisition
(surrogate: ax.models.torch.botorch_modular.surrogate.Surrogate, search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, botorch_acqf_class: Type[botorch.acquisition.acquisition.AcquisitionFunction], options: Optional[Dict[str, Any]] = None, pending_observations: Optional[List[torch.Tensor]] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, objective_thresholds: Optional[torch.Tensor] = None)[source]¶ Bases:
ax.models.torch.botorch_modular.acquisition.Acquisition

acqf
: AcquisitionFunction¶

compute_model_dependencies
(surrogate: ax.models.torch.botorch_modular.surrogate.Surrogate, search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, pending_observations: Optional[List[torch.Tensor]] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Any]] = None) → Dict[str, Any][source]¶ Computes inputs to acquisition function class based on the given surrogate model.
NOTE: When subclassing Acquisition from a superclass where this method returns a nonempty dictionary of kwargs to AcquisitionFunction, call super().compute_model_dependencies and then update that dictionary of options with the options for the subclass you are creating (unless the superclass’ model dependencies should not be propagated to the subclass). See MultiFidelityKnowledgeGradient.compute_model_dependencies for an example.
 Parameters
surrogate – The surrogate object containing the BoTorch Model, with which this Acquisition is to be used.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
pending_observations – A list of tensors, each of which contains points whose evaluation is pending (i.e. that have been submitted for evaluation) for a given outcome. A list of m (k_i x d) feature tensors X for m outcomes and k_i, pending observations for outcome i.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b. (Not used by single task models)
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
options – The options kwarg dict, passed on initialization of the Acquisition object.
 Returns: A dictionary of surrogate modeldependent options, to be passed
as kwargs to BoTorch`AcquisitionFunction` constructor.

surrogate
: Surrogate¶

ax.models.torch.botorch_modular.surrogate module¶

class
ax.models.torch.botorch_modular.surrogate.
Surrogate
(botorch_model_class: Type[botorch.models.model.Model], model_options: Optional[Dict[str, Any]] = None, mll_class: Type[gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood] = <class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>, kernel_class: Optional[Type[gpytorch.kernels.kernel.Kernel]] = None, likelihood: Optional[Type[gpytorch.likelihoods.likelihood.Likelihood]] = None, mll_options: Optional[Dict[str, Any]] = None, kernel_options: Optional[Dict[str, Any]] = None)[source]¶ Bases:
ax.utils.common.base.Base
All classes in ‘botorch_modular’ directory are under construction, incomplete, and should be treated as alpha versions only.
Ax wrapper for BoTorch
Model
, subcomponent ofBoTorchModel
and is not meant to be used outside of it. Parameters
botorch_model_class –
Model
class to be used as the underlying BoTorch model.mll_class –
MarginalLogLikelihood
class to use for modelfitting.model_options – Dictionary of options / kwargs for the BoTorch
Model
constructed duringSurrogate.fit
.kernel_class –
Kernel
class, not yet used. Will be used to construct custom BoTorchModel
in the future.kernel_options – Kernel kwargs, not yet used. Will be used to construct custom BoTorch
Model
in the future.likelihood –
Likelihood
class, not yet used. Will be used to construct custom BoTorchModel
in the future.

best_in_sample_point
(search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: Optional[torch.Tensor], outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Tuple[torch.Tensor, float][source]¶ Finds the best observed point and the corresponding observed outcome values.

best_out_of_sample_point
(search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Tuple[torch.Tensor, torch.Tensor][source]¶ Finds the best predicted point and the corresponding value of the appropriate best point acquisition function.

botorch_model_class
: Type[botorch.models.model.Model]¶

clone_reset
() → ax.models.torch.botorch_modular.surrogate.Surrogate[source]¶

compute_diagnostics
() → Dict[str, Any][source]¶ Computes model diagnostics like crossvalidation measure of fit, etc.

construct
(training_data: botorch.utils.containers.TrainingData, **kwargs: Any) → None[source]¶ Constructs the underlying BoTorch
Model
using the training data. Parameters
training_data – Training data for the model (for one outcome for the default Surrogate, with the exception of batched multioutput case, where training data is formatted with just one X and concatenated Ys).
**kwargs – Optional keyword arguments, expects any of:  “fidelity_features”: Indices of columns in X that represent fidelity.

property
device
¶

property
dtype
¶

fit
(training_data: botorch.utils.containers.TrainingData, search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, state_dict: Optional[Dict[str, torch.Tensor]] = None, refit: bool = True) → None[source]¶ Fits the underlying BoTorch
Model
tom
outcomes.NOTE:
state_dict
andrefit
keyword arguments control how the undelying BoTorchModel
will be fit: whether its parameters will be reoptimized and whether it will be warmstarted from a given state.There are three possibilities:
fit(state_dict=None)
: fit model from stratch (optimize model parameters and set its training data used for inference),fit(state_dict=some_state_dict, refit=True)
: warmstart refit with a state dict of parameters (still reoptimize model parameters and set the training data),fit(state_dict=some_state_dict, refit=False)
: load model parameters without refitting, but set new training data (used in crossvalidation, for example).
 Parameters
data (training) – BoTorch
TrainingData
container with Xs, Ys, and possibly Yvars, to be passed toModel.construct_inputs
in BoTorch.search_space_digest – A SearchSpaceDigest object containing metadata on the features in the trainig data.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.
state_dict – Optional state dict to load.
refit – Whether to reoptimize model parameters.

classmethod
from_botorch
(model: botorch.models.model.Model, mll_class: Type[gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood] = <class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>) → ax.models.torch.botorch_modular.surrogate.Surrogate[source]¶ Instantiate a Surrogate from a preinstantiated Botorch Model.

kernel_class
: Optional[Type[gpytorch.kernels.kernel.Kernel]] = None¶

mll_class
: Type[gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood]¶

property
model
¶

pareto_frontier
() → Tuple[torch.Tensor, torch.Tensor][source]¶ For multiobjective optimization, retrieve Pareto frontier instead of best point.
 Returns: A twotuple of:
tensor of points in the feature space,
tensor of corresponding (multiple) outcomes.

predict
(X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predicts outcomes given a model and input tensor.
 Parameters
model – A botorch Model.
X – A
n x d
tensor of input parameters.
 Returns
The predicted posterior mean as an
n x o
dim tensor. Tensor: The predicted posterior covariance as an x o x o
dim tensor. Return type
Tensor

property
training_data
¶

property
training_data_per_outcome
¶

update
(training_data: botorch.utils.containers.TrainingData, search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, state_dict: Optional[Dict[str, torch.Tensor]] = None, refit: bool = True) → None[source]¶ Updates the surrogate model with new data. In the base
Surrogate
, just callsfit
after checking that this surrogate was not created viaSurrogate.from_botorch
(in which case theModel
comes premade, constructed manually and then supplied toSurrogate
).NOTE: Expects training_data to be all available data, not just the new data since the last time the model was updated.
 Parameters
training_data – Surrogate training_data containing all the data the model should use for inference.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the training data.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.
state_dict – Optional state dict to load.
refit – Whether to reoptimize model parameters or just set the training data used for interence to new training data.
ax.models.torch.botorch_modular.utils module¶

ax.models.torch.botorch_modular.utils.
choose_botorch_acqf_class
(pending_observations: Optional[List[torch.Tensor]] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, objective_thresholds: Optional[torch.Tensor] = None, objective_weights: Optional[torch.Tensor] = None) → Type[botorch.acquisition.acquisition.AcquisitionFunction][source]¶ Chooses a BoTorch AcquisitionFunction class.

ax.models.torch.botorch_modular.utils.
choose_model_class
(Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest) → Type[botorch.models.model.Model][source]¶ Chooses a BoTorch Model using the given data (currently just Yvars) and its properties (information about task and fidelity features).
 Parameters
Yvars – List of tensors, each representing observation noise for a given outcome, where outcomes are in the same order as in Xs.
task_features – List of columns of X that are tasks.
fidelity_features – List of columns of X that are fidelity parameters.
 Returns
A BoTorch Model class.

ax.models.torch.botorch_modular.utils.
construct_acquisition_and_optimizer_options
(acqf_options: Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]], model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Tuple[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]], Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]][source]¶ Extract acquisition and optimizer options from model_gen_options.
ax.models.torch.cbo_lcea module¶

class
ax.models.torch.cbo_lcea.
LCEABO
(decomposition: Dict[str, List[str]], cat_feature_dict: Optional[Dict] = None, embs_feature_dict: Optional[Dict] = None, context_weight_dict: Optional[Dict] = None, embs_dim_list: Optional[List[int]] = None, gp_model_args: Optional[Dict[str, Any]] = None)[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
Does Bayesian optimization with Latent Context Embedding Additive (LCEA) GP. The parameter space decomposition must be provided.
 Parameters
decomposition – Keys are context names. Values are the lists of parameter names belong to the context, e.g. {‘context1’: [‘p1_c1’, ‘p2_c1’],’context2’: [‘p1_c2’, ‘p2_c2’]}.
gp_model_args – Dictionary of kwargs to pass to GP model training.  train_embedding: Boolen. If true, we will train context embedding; otherwise, we use pretrained embeddings from embds_feature_dict only. Default is True.

Xs
: List[Tensor]¶

Ys
: List[Tensor]¶

Yvars
: List[Tensor]¶

best_point
(bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ Identify the current best point, satisfying the constraints in the same format as to gen.
Return None if no such point can be identified.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
dtensor of the best point.

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

fit
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

get_and_fit_model
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, torch.Tensor]] = None, fidelity_model_id: Optional[int] = None, **kwargs: Any) → botorch.models.gpytorch.GPyTorchModel[source]¶ Get a fitted LCEAGP model for each outcome. :param Xs: X for each outcome. :param Ys: Y for each outcome. :param Yvars: Noise variance of Y for each outcome.
Returns: Fitted LCEAGP model.

model
: Optional[Model]¶

ax.models.torch.cbo_lcea.
get_map_model
(train_X: torch.Tensor, train_Y: torch.Tensor, train_Yvar: torch.Tensor, decomposition: Dict[str, List[int]], train_embedding: bool = True, cat_feature_dict: Optional[Dict] = None, embs_feature_dict: Optional[Dict] = None, embs_dim_list: Optional[List[int]] = None, context_weight_dict: Optional[Dict] = None) → Tuple[botorch.models.contextual.LCEAGP, gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood][source]¶ Obtain MAP fitting of Latent Context Embedding Additive (LCEA) GP.
ax.models.torch.cbo_lcem module¶

class
ax.models.torch.cbo_lcem.
LCEMBO
(context_cat_feature: Optional[torch.Tensor] = None, context_emb_feature: Optional[torch.Tensor] = None, embs_dim_list: Optional[List[int]] = None)[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
Does Bayesian optimization with LCEM GP.

Xs
: List[Tensor]¶

Ys
: List[Tensor]¶

Yvars
: List[Tensor]¶

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

get_and_fit_model
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, torch.Tensor]] = None, fidelity_model_id: Optional[int] = None, **kwargs: Any) → botorch.models.model_list_gp_regression.ModelListGP[source]¶ Get a fitted multitask contextual GP model for each outcome. :param Xs: List of X data, one tensor per outcome. :param Ys: List of Y data, one tensor per outcome. :param Yvars: List of Noise variance of Yvar data, one tensor per outcome. :param task_features: List of columns of X that are tasks.
Returns: ModeListGP that each model is a fitted LCEM GP model.

model
: Optional[Model]¶

ax.models.torch.cbo_sac module¶

class
ax.models.torch.cbo_sac.
SACBO
(decomposition: Dict[str, List[str]])[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
Does Bayesian optimization with structural additive contextual GP (SACGP). The parameter space decomposition must be provided.
 Parameters
decomposition – Keys are context names. Values are the lists of parameter names belong to the context, e.g. {‘context1’: [‘p1_c1’, ‘p2_c1’],’context2’: [‘p1_c2’, ‘p2_c2’]}.

Xs
: List[Tensor]¶

Ys
: List[Tensor]¶

Yvars
: List[Tensor]¶

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

fit
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

get_and_fit_model
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, torch.Tensor]] = None, fidelity_model_id: Optional[int] = None, **kwargs: Any) → botorch.models.gpytorch.GPyTorchModel[source]¶ Get a fitted StructuralAdditiveContextualGP model for each outcome. :param Xs: X for each outcome. :param Ys: Y for each outcome. :param Yvars: Noise variance of Y for each outcome.
Returns: Fitted StructuralAdditiveContextualGP model.

model
: Optional[Model]¶
ax.models.torch.frontier_utils module¶

ax.models.torch.frontier_utils.
get_default_frontier_evaluator
() → Callable[[ax.models.torch_base.TorchModel, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]][source]¶

ax.models.torch.frontier_utils.
get_weighted_mc_objective_and_objective_thresholds
(objective_weights: torch.Tensor, objective_thresholds: torch.Tensor) → Tuple[botorch.acquisition.multi_objective.objective.WeightedMCMultiOutputObjective, torch.Tensor][source]¶ Construct weighted objective and apply the weights to objective thresholds.
 Parameters
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
objective_thresholds – A tensor containing thresholds forming a reference point from which to calculate pareto frontier hypervolume. Points that do not dominate the objective_thresholds contribute nothing to hypervolume.
 Returns
The objective
The objective thresholds
 Return type
A twoelement tuple with the objective and objective thresholds
ax.models.torch.fully_bayesian module¶
Models and utilities for fully bayesian inference.
TODO: move some of this into botorch.
References
 Eriksson2021saasbo
D. Eriksson, M. Jankowiak. HighDimensional Bayesian Optimization with Sparse AxisAligned Subspaces. Proceedings of the Thirty Seventh Conference on Uncertainty in Artificial Intelligence, 2021.
 Eriksson2021nas
D. Eriksson, P. Chuang, S. Daulton, et al. LatencyAware Neural Architecture Search with MultiObjective Bayesian Optimization. ICML AutoML Workshop, 2021.

class
ax.models.torch.fully_bayesian.
FullyBayesianBotorchModel
(model_constructor: Callable[[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[int], List[int], List[str], Optional[Dict[str, torch.Tensor]], Any], botorch.models.model.Model] = <function get_and_fit_model_mcmc>, model_predictor: Callable[[botorch.models.model.Model, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]] = <function predict_from_model_mcmc>, acqf_constructor: Callable[[botorch.models.model.Model, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], Any], botorch.acquisition.acquisition.AcquisitionFunction] = <function get_fully_bayesian_acqf>, acqf_optimizer: Callable[[botorch.acquisition.acquisition.AcquisitionFunction, torch.Tensor, int, Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[Dict[int, float]], Optional[Callable[[torch.Tensor], torch.Tensor]], Any], Tuple[torch.Tensor, torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: Callable[[ax.models.torch_base.TorchModel, List[Tuple[float, float]], torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Dict[int, float]], Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]], Optional[Dict[int, float]]], Optional[torch.Tensor]] = <function recommend_best_observed_point>, refit_on_cv: bool = False, refit_on_update: bool = True, warm_start_refitting: bool = True, use_input_warping: bool = False, use_saas: Optional[bool] = None, num_samples: int = 512, warmup_steps: int = 1024, thinning: int = 16, max_tree_depth: int = 6, disable_progbar: bool = False, gp_kernel: str = 'matern', verbose: bool = False, **kwargs: Any)[source]¶ Bases:
ax.models.torch.fully_bayesian.FullyBayesianBotorchModelMixin
,ax.models.torch.botorch.BotorchModel
Fully Bayesian Model that uses NUTS to sample from hyperparameter posterior.
This includes support for using sparse axisaligned subspace priors (SAAS). See [Eriksson2021saasbo] for details.

class
ax.models.torch.fully_bayesian.
FullyBayesianBotorchModelMixin
[source]¶ Bases:
object

model
: Optional[botorch.models.model.Model] = None¶


class
ax.models.torch.fully_bayesian.
FullyBayesianMOOBotorchModel
(model_constructor: Callable[[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[int], List[int], List[str], Optional[Dict[str, torch.Tensor]], Any], botorch.models.model.Model] = <function get_and_fit_model_mcmc>, model_predictor: Callable[[botorch.models.model.Model, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]] = <function predict_from_model_mcmc>, acqf_constructor: Callable[[botorch.models.model.Model, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], Any], botorch.acquisition.acquisition.AcquisitionFunction] = <function get_fully_bayesian_acqf_nehvi>, acqf_optimizer: Callable[[botorch.acquisition.acquisition.AcquisitionFunction, torch.Tensor, int, Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[Dict[int, float]], Optional[Callable[[torch.Tensor], torch.Tensor]], Any], Tuple[torch.Tensor, torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: Callable[[ax.models.torch_base.TorchModel, List[Tuple[float, float]], torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Dict[int, float]], Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]], Optional[Dict[int, float]]], Optional[torch.Tensor]] = <function recommend_best_observed_point>, frontier_evaluator: Callable[[ax.models.torch_base.TorchModel, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = <function pareto_frontier_evaluator>, refit_on_cv: bool = False, refit_on_update: bool = True, warm_start_refitting: bool = False, use_input_warping: bool = False, num_samples: int = 512, warmup_steps: int = 1024, thinning: int = 16, max_tree_depth: int = 6, use_saas: Optional[bool] = None, disable_progbar: bool = False, gp_kernel: str = 'matern', verbose: bool = False, **kwargs: Any)[source]¶ Bases:
ax.models.torch.fully_bayesian.FullyBayesianBotorchModelMixin
,ax.models.torch.botorch_moo.MultiObjectiveBotorchModel
Fully Bayesian Model that uses qNEHVI.
This includes support for using qNEHVI + SAASBO as in [Eriksson2021nas].

ax.models.torch.fully_bayesian.
compute_dists
(X: torch.Tensor, Z: torch.Tensor, lengthscale: torch.Tensor) → torch.Tensor[source]¶ Compute kernel distances.
TODO: use gpytorch Distance module. This will require some care to make sure jit compilation works as expected.

ax.models.torch.fully_bayesian.
get_and_fit_model_mcmc
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, torch.Tensor]] = None, refit_model: bool = True, use_input_warping: bool = False, use_loocv_pseudo_likelihood: bool = False, num_samples: int = 512, warmup_steps: int = 1024, thinning: int = 16, max_tree_depth: int = 6, disable_progbar: bool = False, gp_kernel: str = 'matern', verbose: bool = False, **kwargs: Any) → botorch.models.gpytorch.GPyTorchModel[source]¶ Instantiates a batched GPyTorchModel(ModelListGP) based on the given data and fit the model based on MCMC in pyro. The batch dimension corresponds to sampled hyperparameters from MCMC.

ax.models.torch.fully_bayesian.
get_fully_bayesian_acqf
(model: botorch.models.model.Model, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None, X_pending: Optional[torch.Tensor] = None, acqf_constructor: Callable[[botorch.models.model.Model, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], Any], botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEI>, **kwargs: Any) → botorch.acquisition.acquisition.AcquisitionFunction[source]¶

ax.models.torch.fully_bayesian.
get_fully_bayesian_acqf_nehvi
(model: botorch.models.model.Model, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None, X_pending: Optional[torch.Tensor] = None, **kwargs: Any) → botorch.acquisition.acquisition.AcquisitionFunction[source]¶

ax.models.torch.fully_bayesian.
matern_kernel
(X: torch.Tensor, Z: torch.Tensor, lengthscale: torch.Tensor, nu: float = 2.5) → torch.Tensor[source]¶ Scaled Matern kernel.

ax.models.torch.fully_bayesian.
predict_from_model_mcmc
(model: botorch.models.model.Model, X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predicts outcomes given a model and input tensor.
This method integrates over the hyperparameter posterior.
 Parameters
model – A batched botorch Model where the batch dimension corresponds to sampled hyperparameters.
X – A n x d tensor of input parameters.
 Returns
The predicted posterior mean as an n x odim tensor. Tensor: The predicted posterior covariance as a n x o x odim tensor.
 Return type
Tensor

ax.models.torch.fully_bayesian.
rbf_kernel
(X: torch.Tensor, Z: torch.Tensor, lengthscale: torch.Tensor) → torch.Tensor[source]¶ Scaled RBF kernel.

ax.models.torch.fully_bayesian.
run_inference
(pyro_model: Callable, X: torch.Tensor, Y: torch.Tensor, Yvar: torch.Tensor, num_samples: int = 512, warmup_steps: int = 1024, thinning: int = 16, use_input_warping: bool = False, max_tree_depth: int = 6, disable_progbar: bool = False, gp_kernel: str = 'matern', verbose: bool = False, task_feature: Optional[int] = None, rank: Optional[int] = None) → Dict[str, torch.Tensor][source]¶

ax.models.torch.fully_bayesian.
single_task_pyro_model
(X: torch.Tensor, Y: torch.Tensor, Yvar: torch.Tensor, use_input_warping: bool = False, eps: float = 1e07, gp_kernel: str = 'matern', task_feature: Optional[int] = None, rank: Optional[int] = None) → None[source]¶ Instantiates a single task pyro model for running fully bayesian inference.
 Parameters
X – A n x d tensor of input parameters.
Y – A n x 1 tensor of output.
Yvar – A n x 1 tensor of observed noise.
use_input_warping – A boolean indicating whether to use input warping
task_feature – Column index of task feature in X.
gp_kernel – kernel name. Currently only two kernels are supported: “matern” for Matern Kernel and “rbf” for RBFKernel.
rank – num of latent task features to learn for task covariance.
ax.models.torch.fully_bayesian_model_utils module¶

ax.models.torch.fully_bayesian_model_utils.
load_mcmc_samples_to_model
(model: botorch.models.gpytorch.GPyTorchModel, mcmc_samples: Dict) → None[source]¶ Load MCMC samples into GPyTorchModel.

ax.models.torch.fully_bayesian_model_utils.
load_pyro
(func: Callable) → Callable[source]¶ A decorator to import pyro

ax.models.torch.fully_bayesian_model_utils.
pyro_sample_input_warping
(dim: int, pyro: Any = None, **tkwargs: Any) → Tuple[torch.Tensor, torch.Tensor][source]¶

ax.models.torch.fully_bayesian_model_utils.
pyro_sample_mean
(pyro: Any = None, **tkwargs: Any) → torch.Tensor[source]¶

ax.models.torch.fully_bayesian_model_utils.
pyro_sample_noise
(pyro: Any = None, **tkwargs: Any) → torch.Tensor[source]¶
ax.models.torch.posterior_mean module¶

ax.models.torch.posterior_mean.
get_PosteriorMean
(model: botorch.models.model.Model, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, X_observed: Optional[torch.Tensor] = None, X_pending: Optional[torch.Tensor] = None, **kwargs: Any) → botorch.acquisition.acquisition.AcquisitionFunction[source]¶ Instantiates a PosteriorMean acquisition function.
Note: If no OutcomeConstraints given, return an analytic acquisition function. This requires {optimizer_kwargs: {joint_optimization: True}} or an optimizer that does not assume pending point support.
 Parameters
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
X_observed – A tensor containing points observed for all objective outcomes and outcomes that appear in the outcome constraints (if there are any).
X_pending – A tensor containing points whose evaluation is pending (i.e. that have been submitted for evaluation) present for all objective outcomes and outcomes that appear in the outcome constraints (if there are any).
 Returns
The instantiated acquisition function.
 Return type
PosteriorMean
ax.models.torch.rembo module¶

class
ax.models.torch.rembo.
REMBO
(A: torch.Tensor, initial_X_d: torch.Tensor, bounds_d: List[Tuple[float, float]], **kwargs: Any)[source]¶ Bases:
ax.models.torch.botorch.BotorchModel
Implements REMBO (Bayesian optimization in a linear subspace).
The (D x d) projection matrix A must be provided, and must be that used for the initialization. In the original REMBO paper A ~ N(0, 1). Box bounds in the lowd space must also be provided, which in the REMBO paper should be [(sqrt(d), sqrt(d)]^d.
Function evaluations happen in the highD space, and so the arms on the experiment will also be tracked in the highD space. This class maintains a list of points in the lowd spac that have been launched, so we can match arms in highD space back to their lowd point on update.
 Parameters
A – (D x d) projection matrix.
initial_X_d – Points in lowd space for initial data.
bounds_d – Box bounds in the lowd space.
kwargs – kwargs for BotorchModel init

Xs
: List[Tensor]¶

Ys
: List[Tensor]¶

Yvars
: List[Tensor]¶

best_point
(bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[torch.Tensor][source]¶ Identify the current best point, satisfying the constraints in the same format as to gen.
Return None if no such point can be identified.
 Parameters
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value in the best point.
model_gen_options – A config dictionary that can contain modelspecific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
dtensor of the best point.

cross_validate
(Xs_train: List[torch.Tensor], Ys_train: List[torch.Tensor], Yvars_train: List[torch.Tensor], X_test: torch.Tensor, **kwargs: Any) → Tuple[torch.Tensor, torch.Tensor][source]¶ Do cross validation with the given training and test sets.
Training set is given in the same format as to fit. Test set is given in the same format as to predict.
 Parameters
Xs_train – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys_train – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars_train – The variances of each entry in Ys, same shape.
X_test – (j x d) tensor of the j points at which to make predictions.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

device
: Optional[torch.device]¶

dtype
: Optional[torch.dtype]¶

fit
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], search_space_digest: ax.core.search_space.SearchSpaceDigest, metric_names: List[str], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶ Fit model to m outcomes.
 Parameters
Xs – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.
Ys – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.
Yvars – The variances of each entry in Ys, same shape.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.

from_01
(X_d01: torch.Tensor) → torch.Tensor[source]¶ Map points from [0, 1] to bounds_d.
 Parameters
X_d01 – Tensor in [0, 1]
Returns: Tensor in bounds_d.

gen
(n: int, bounds: List[Tuple[float, float]], objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, linear_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[torch.Tensor]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None, rounding_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], Optional[List[Optional[Dict[str, Any]]]]][source]¶ Generate new candidates.
 Parameters
n – Number of candidates to generate.
bounds – A list of (lower, upper) tuples for each column of X.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b.
linear_constraints – A tuple of (A, b). For k linear constraints on ddimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
fixed_features – A map {feature_index: value} for features that should be fixed to a particular value during generation.
pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.
model_gen_options – A config dictionary that can contain modelspecific options.
rounding_func – A function that rounds an optimization result appropriately (i.e., according to roundtrip transformations).
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multifidelity optimization.
 Returns
4element tuple containing
(n x d) tensor of generated points.
ntensor of weights for each point.
Generation metadata
 Dictionary of modelspecific metadata for the given
generation candidates

model
: Optional[Model]¶

predict
(X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predict
 Parameters
X – (j x d) tensor of the j points at which to make predictions.
 Returns
2element tuple containing
(j x m) tensor of outcome predictions at X.
(j x m x m) tensor of predictive covariances at X. cov[j, m1, m2] is Cov[[email protected], [email protected]].

project_down
(X_D: torch.Tensor) → torch.Tensor[source]¶ Map points in the highD space to the lowd space by looking them up in self.X_d.
We assume that X_D = self.project_up(self.X_d), except possibly with rows shuffled. If a value in X_d cannot be found for each row in X_D, an error will be raised.
This is quite fast relative to model fitting, so we do it in O(n^2) time and don’t worry about it.
 Parameters
X_D – Tensor in highD space.
 Returns
Tensor in lowd space.
 Return type
X_d

to_01
(X_d: torch.Tensor) → torch.Tensor[source]¶ Map points from bounds_d to [0, 1].
 Parameters
X_d – Tensor in bounds_d
Returns: Tensor in [0, 1].

update
(Xs: List[torch.Tensor], Ys: List[torch.Tensor], Yvars: List[torch.Tensor], candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, **kwargs: Any) → None[source]¶ Update the model.
Updating the model requires both existing and additional data. The data passed into this method will become the new training data.
 Parameters
Xs – Existing + additional data for the model, in the same format as for fit.
Ys – Existing + additional data for the model, in the same format as for fit.
Yvars – Existing + additional data for the model, in the same format as for fit.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
metric_names – Names of each outcome Y in Ys.
candidate_metadata – Modelproduced metadata for candidates, in the order corresponding to the Xs.
ax.models.torch.utils module¶

class
ax.models.torch.utils.
SubsetModelData
(model: botorch.models.model.Model, objective_weights: torch.Tensor, outcome_constraints: Union[Tuple[torch.Tensor, torch.Tensor], NoneType], objective_thresholds: Union[torch.Tensor, NoneType], indices: torch.Tensor)[source]¶ Bases:
object

indices
: torch.Tensor¶

model
: botorch.models.model.Model¶

objective_thresholds
: Optional[torch.Tensor]¶

objective_weights
: torch.Tensor¶

outcome_constraints
: Optional[Tuple[torch.Tensor, torch.Tensor]]¶


ax.models.torch.utils.
get_botorch_objective
(model: botorch.models.model.Model, objective_weights: torch.Tensor, use_scalarized_objective: bool = True, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, objective_thresholds: Optional[torch.Tensor] = None, X_observed: Optional[torch.Tensor] = None) → botorch.acquisition.objective.AcquisitionObjective[source]¶ Constructs a BoTorch AcquisitionObjective object.
 Parameters
model – A BoTorch Model
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
use_scalarized_objective – A boolean parameter that defaults to True, specifying whether ScalarizedObjective should be used. NOTE: when using outcome_constraints, use_scalarized_objective will be ignored.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
objective_thresholds – A tensor containing thresholds forming a reference point from which to calculate pareto frontier hypervolume. Points that do not dominate the objective_thresholds contribute nothing to hypervolume.
X_observed – Observed points that are feasible and appear in the objective or the constraints. None if there are no such points.
 Returns
ScalarizedObjective, LinearMCOObjective, ConstrainedMCObjective.
 Return type
A BoTorch AcquisitionObjective object. It will be one of

ax.models.torch.utils.
get_out_of_sample_best_point_acqf
(model: botorch.models.model.Model, Xs: List[torch.Tensor], X_observed: torch.Tensor, objective_weights: torch.Tensor, mc_samples: int = 512, fixed_features: Optional[Dict[int, float]] = None, fidelity_features: Optional[List[int]] = None, target_fidelities: Optional[Dict[int, float]] = None, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, seed_inner: Optional[int] = None, qmc: bool = True, **kwargs: Any) → Tuple[botorch.acquisition.acquisition.AcquisitionFunction, Optional[List[int]]][source]¶ Picks an appropriate acquisition function to find the best outofsample (predicted by the given surrogate model) point and instantiates it.
NOTE: Typically the appropriate function is the posterior mean, but can differ to account for fidelities etc.

ax.models.torch.utils.
is_noiseless
(model: botorch.models.model.Model) → bool[source]¶ Check if a given (singletask) botorch model is noiseless

ax.models.torch.utils.
normalize_indices
(indices: List[int], d: int) → List[int][source]¶ Normalize a list of indices to ensure that they are positive.
 Parameters
indices – A list of indices (may contain negative indices for indexing “from the back”).
d – The dimension of the tensor to index.
 Returns
A normalized list of indices such that each index is between 0 and d1.

ax.models.torch.utils.
pick_best_out_of_sample_point_acqf_class
(outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, mc_samples: int = 512, qmc: bool = True, seed_inner: Optional[int] = None) → Tuple[Type[botorch.acquisition.acquisition.AcquisitionFunction], Dict[str, Any]][source]¶

ax.models.torch.utils.
predict_from_model
(model: botorch.models.model.Model, X: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]¶ Predicts outcomes given a model and input tensor.
 Parameters
model – A botorch Model.
X – A n x d tensor of input parameters.
 Returns
The predicted posterior mean as an n x odim tensor. Tensor: The predicted posterior covariance as a n x o x odim tensor.
 Return type
Tensor

ax.models.torch.utils.
randomize_objective_weights
(objective_weights: torch.Tensor, **acquisition_function_kwargs: Any) → torch.Tensor[source]¶ Generate a random weighting based on acquisition function settings.
 Parameters
objective_weights – Base weights to multiply by random values..
**acquisition_function_kwargs – Kwargs containing weight generation algorithm options.
 Returns
A normalized list of indices such that each index is between 0 and d1.

ax.models.torch.utils.
subset_model
(model: botorch.models.model.Model, objective_weights: torch.Tensor, outcome_constraints: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, objective_thresholds: Optional[torch.Tensor] = None) → ax.models.torch.utils.SubsetModelData[source]¶ Subset a botorch model to the outputs used in the optimization.
 Parameters
model – A BoTorch Model. If the model does not implement the subset_outputs method, this function is a nullop and returns the input arguments.
objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.
objective_thresholds – The mdim tensor of objective thresholds. There is one for each modeled metric.
outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)
 Returns
A SubsetModelData dataclass containing the model, objective_weights, outcome_constraints, objective thresholds, all subset to only those outputs that appear in either the objective weights or the outcome constraints, along with the indices of the outputs.