ax.models¶

Base Models & Utilities¶

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.

classmethod deserialize_state(serialized_state: Dict[str, Any]) → Dict[str, Any][source]¶: Restores model’s state from its serialized form, to the format it expects to receive as kwargs.

feature_importances() → Any[source]¶

classmethod serialize_state(raw_state: Dict[str, Any]) → Dict[str, Any][source]¶: Serialized output of self._get_state to a JSON-ready dict. This may involve storing part of state in files / external storage and saving handles for that storage in the resulting serialized state.

ax.models.discrete_base module¶

class ax.models.discrete_base.DiscreteModel[source]¶

Bases: 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[ndarray], outcome_constraints: Optional[Tuple[ndarray, 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, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = 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[ndarray, 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:

2-element 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[m1@j, m2@j].

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[ndarray], outcome_constraints: Optional[Tuple[ndarray, 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, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = 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 model-specific options.

Returns:

2-element 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[ndarray, ndarray][source]¶

Predict

Parameters:

X – List of the j parameterizations at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

ax.models.torch_base module¶

class ax.models.torch_base.TorchGenResults(points: ~torch.Tensor, weights: ~torch.Tensor, gen_metadata: ~typing.Dict[str, ~typing.Any] = <factory>, candidate_metadata: ~typing.Optional[~typing.List[~typing.Optional[~typing.Dict[str, ~typing.Any]]]] = None)[source]¶

Bases: object

points: (n x d) Tensor of generated points. weights: n-tensor of weights for each point. gen_metadata: Generation metadata Dictionary of model-specific metadata for the given

generation candidates

candidate_metadata: Optional[List[Optional[Dict[str, Any]]]] = None¶

gen_metadata: Dict[str, Any]¶

points: Tensor¶

weights: Tensor¶

class ax.models.torch_base.TorchModel[source]¶

Bases: Model

This class specifies the interface for a torch-based model.

These methods should be implemented to have access to all of the features of Ax.

best_point(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → Optional[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:

search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

d-tensor of the best point.

cross_validate(datasets: List[SupervisedDataset], metric_names: List[str], X_test: Tensor, search_space_digest: SearchSpaceDigest) → Tuple[Tensor, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
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.

Returns:

2-element 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[m1@j, m2@j].

device: Optional[device] = None¶

dtype: Optional[dtype] = None¶

evaluate_acquisition_function(X: Tensor, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, acq_options: Optional[Dict[str, Any]] = None) → 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.
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).
acq_options – Keyword arguments used to contruct the acquisition function.

Returns:

A single-element tensor with the acquisition value for these points.

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

gen(n: int, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[source]¶

Generate new candidates.

Parameters:

n – Number of candidates to generate.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

A TorchGenResult container.

predict(X: Tensor) → Tuple[Tensor, Tensor][source]¶

Predict

Parameters:

X – (j x d) tensor of the j points at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

update(datasets: List[Optional[SupervisedDataset]], metric_names: List[str], search_space_digest: SearchSpaceDigest, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome). None means that there is no additional data for the corresponding outcome.
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

class ax.models.torch_base.TorchOptConfig(objective_weights: ~torch.Tensor, outcome_constraints: ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = None, objective_thresholds: ~typing.Optional[~torch.Tensor] = None, linear_constraints: ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = None, fixed_features: ~typing.Optional[~typing.Dict[int, float]] = None, pending_observations: ~typing.Optional[~typing.List[~torch.Tensor]] = None, model_gen_options: ~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]] = <factory>, rounding_func: ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]] = None, opt_config_metrics: ~typing.Optional[~typing.Dict[str, ~ax.core.metric.Metric]] = None, is_moo: bool = False)[source]¶

Bases: object

Container for lightweight representation of optimization arguments.

This is used for communicating between modelbridge and models. This is an ephemeral object and not meant to be stored / serialized.

objective_weights¶

If doing multi-objective optimization, these denote which objectives should be maximized and which should be minimized. Otherwise, the objective is to maximize a weighted sum of the columns of f(x). These are the weights.

Type:: torch.Tensor

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.

Type:: Optional[Tuple[torch.Tensor, torch.Tensor]]

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.

Type:: Optional[torch.Tensor]

linear_constraints¶

A tuple of (A, b). For k linear constraints on d-dimensional x, A is (k x d) and b is (k x 1) such that A x <= b for feasible x.

Type:: Optional[Tuple[torch.Tensor, torch.Tensor]]

fixed_features¶

A map {feature_index: value} for features that should be fixed to a particular value during generation.

Type:: Optional[Dict[int, float]]

pending_observations¶

A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.

Type:: Optional[List[torch.Tensor]]

model_gen_options¶

A config dictionary that can contain model-specific options.

Type:: Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any], ax.core.optimization_config.OptimizationConfig]]]

rounding_func¶

A function that rounds an optimization result appropriately (i.e., according to round-trip transformations).

Type:: Optional[Callable[[torch.Tensor], torch.Tensor]]

opt_config_metrics¶

A dictionary of metrics that are included in the optimization config.

Type:: Optional[Dict[str, ax.core.metric.Metric]]

is_moo¶

A boolean denoting whether this is for an MOO problem.

Type:: bool

fixed_features: Optional[Dict[int, float]] = None¶

is_moo: bool = False¶

linear_constraints: Optional[Tuple[Tensor, Tensor]] = None¶

model_gen_options: Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]¶

objective_thresholds: Optional[Tensor] = None¶

objective_weights: Tensor¶

opt_config_metrics: Optional[Dict[str, Metric]] = None¶

outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None¶

pending_observations: Optional[List[Tensor]] = None¶

rounding_func: Optional[Callable[[Tensor], Tensor]] = None¶

ax.models.model_utils module¶

ax.models.model_utils.add_fixed_features(tunable_points: ndarray, d: int, fixed_features: Optional[Dict[int, float]], tunable_feature_indices: ndarray) → 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 d-dimensional space, defined by bounds.

Return type:

points

ax.models.model_utils.as_array(x: Union[Tensor, ndarray, Tuple[Union[Tensor, ndarray], ...]]) → Union[ndarray, Tuple[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[Tensor], List[ndarray]], model: TorchModel, bounds: List[Tuple[float, float]], objective_weights: Optional[Union[Tensor, ndarray]], outcome_constraints: Optional[Tuple[Union[Tensor, ndarray], Union[Tensor, ndarray]]] = None, linear_constraints: Optional[Tuple[Union[Tensor, ndarray], Union[Tensor, ndarray]]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None) → Optional[Tuple[Union[Tensor, 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 best-objective 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 d-dimensional 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:

d-array of the best point,
utility at the best point.

Return type:

A two-element tuple or None if no feasible point exist. In tuple

ax.models.model_utils.best_observed_point(model: TorchModel, bounds: List[Tuple[float, float]], objective_weights: Optional[Union[Tensor, ndarray]], outcome_constraints: Optional[Tuple[Union[Tensor, ndarray], Union[Tensor, ndarray]]] = None, linear_constraints: Optional[Tuple[Union[Tensor, ndarray], Union[Tensor, ndarray]]] = None, fixed_features: Optional[Dict[int, float]] = None, options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None) → Optional[Union[Tensor, 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 best-objective 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 d-dimensional 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 d-array of the best point, or None if no feasible point exists.

ax.models.model_utils.check_duplicate(point: ndarray, points: 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[ndarray, ndarray], point: ndarray) → Tuple[bool, ndarray][source]¶

Check if a point satisfies parameter constraints.

Parameters:

linear_constraints – A tuple of (A, b). For k linear constraints on d-dimensional x, A is (k x d) and b is (k x 1) such that A x <= b.
point – A candidate point in d-dimensional space, as a (1 x d) matrix.

Returns:

2-element 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[float, int]]][source]¶

ax.models.model_utils.filter_constraints_and_fixed_features(X: Union[Tensor, ndarray], bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[Union[Tensor, ndarray], Union[Tensor, ndarray]]] = None, fixed_features: Optional[Dict[int, float]] = None) → Union[Tensor, 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 d-dimensional 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[Tensor], List[ndarray]], objective_weights: Union[Tensor, ndarray], outcome_constraints: Optional[Tuple[Union[Tensor, ndarray], Union[Tensor, ndarray]]] = None) → Union[Tensor, 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: 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, ndarray, Optional[Dict[int, float]]], ndarray], n: int, d: int, tunable_feature_indices: ndarray, linear_constraints: Optional[Tuple[ndarray, ndarray]] = None, deduplicate: bool = False, max_draws: Optional[int] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[ndarray], ndarray]] = None, existing_points: Optional[ndarray] = None) → Tuple[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) → 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: 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.types¶

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: ThompsonSampler

Generator for Thompson sampling using Empirical Bayes estimates.

The generator applies positive-part James-Stein 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: 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[ndarray], outcome_constraints: Optional[Tuple[ndarray, 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, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = 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 model-specific options.

Returns:

2-element 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: 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[ndarray], outcome_constraints: Optional[Tuple[ndarray, 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, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = 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 model-specific options.

Returns:

2-element 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[ndarray, ndarray][source]¶

Predict

Parameters:

X – List of the j parameterizations at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

Random Models¶

ax.models.random.base module¶

class ax.models.random.base.RandomModel(deduplicate: bool = True, seed: Optional[int] = None, generated_points: Optional[ndarray] = None, fallback_to_sample_polytope: bool = False)[source]¶

Bases: 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[ndarray, ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None, rounding_func: Optional[Callable[[ndarray], ndarray]] = None) → Tuple[ndarray, 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 d-dimensional 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 round-trip transformations).

Returns:

2-element tuple containing

(n x d) array of generated points.
Uniform weights, an n-array 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: 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.

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[ndarray] = None, fallback_to_sample_polytope: bool = False)[source]¶

Bases: 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: Optional[SobolEngine]¶: Return a singleton SobolEngine.

gen(n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[ndarray, ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None, rounding_func: Optional[Callable[[ndarray], ndarray]] = None) → Tuple[ndarray, 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 d-dimensional 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 round-trip transformations) but unused here.

Returns:

2-element tuple containing

(n x d) array of generated points.
Uniform weights, an n-array of ones for each point.

init_engine(n_tunable_features: int) → SobolEngine[source]¶

Initialize singleton SobolEngine, only on gen.

Parameters:: n_tunable_features – The number of features which can be searched over.
Returns:: SobolEngine, which can generate Sobol points.

ax.models.random.alebo_initializer module¶

class ax.models.random.alebo_initializer.ALEBOInitializer(B: ndarray, nsamp: int = 10000, init_bound: int = 16, seed: Optional[int] = None)[source]¶

Bases: UniformGenerator

Sample in a low-dimensional 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 up-projection 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.
seed – seed for UniformGenerator

gen(n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[ndarray, ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None, rounding_func: Optional[Callable[[ndarray], ndarray]] = None) → Tuple[ndarray, 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 d-dimensional 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 round-trip transformations).

Returns:

2-element tuple containing

(n x d) array of generated points.
Uniform weights, an n-array of ones for each point.

ax.models.random.rembo_initializer module¶

class ax.models.random.rembo_initializer.REMBOInitializer(A: ndarray, bounds_d: List[Tuple[float, float]], seed: Optional[int] = None)[source]¶

Bases: UniformGenerator

Sample in a low-dimensional linear embedding.

Generates points in [-1, 1]^D by generating points in a d-dimensional 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 low-d space
seed – seed for UniformGenerator

gen(n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[ndarray, ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None, rounding_func: Optional[Callable[[ndarray], ndarray]] = None) → Tuple[ndarray, 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 d-dimensional 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 round-trip transformations).

Returns:

2-element tuple containing

(n x d) array of generated points.
Uniform weights, an n-array of ones for each point.

project_up(X: ndarray) → ndarray[source]¶: Project to high-dimensional space.

Torch Models & Utilities¶

ax.models.torch.alebo module¶

class ax.models.torch.alebo.ALEBO(B: Tensor, laplace_nsamp: int = 25, fit_restarts: int = 10)[source]¶

Bases: 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 high-D 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(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → Optional[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:

search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

d-tensor of the best point.

cross_validate(datasets: List[SupervisedDataset], X_test: Tensor, **kwargs: Any) → Tuple[Tensor, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
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.

Returns:

2-element 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[m1@j, m2@j].

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

gen(n: int, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[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[Tensor], Ys: List[Tensor], Yvars: List[Tensor], state_dicts: Optional[List[MutableMapping[str, Tensor]]] = None) → 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: Tensor) → Tuple[Tensor, Tensor][source]¶

Predict

Parameters:

X – (j x d) tensor of the j points at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

update(datasets: List[SupervisedDataset], 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome). None means that there is no additional data for the corresponding outcome.
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

class ax.models.torch.alebo.ALEBOGP(B: Tensor, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor)[source]¶

Bases: 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 non-batch 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 non-batch 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: Tensor, output_indices: Optional[List[int]] = None, observation_noise: Union[bool, Tensor] = False, posterior_transform: Optional[PosteriorTransform] = None, **kwargs: Any) → GPyTorchPosterior[source]¶

Computes the posterior over model outputs at the provided points.

Parameters:

X – A (batch_shape) x q x d-dim 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 multi-output). 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).
posterior_transform – An optional PosteriorTransform.

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: Tensor, batch_shape: Size)[source]¶

Bases: 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.

forward(x1: Tensor, x2: Tensor, diag: bool = False, last_dim_is_batch: bool = False, **params: Any) → Tensor[source]¶: Compute kernel distance.

training: bool¶

ax.models.torch.alebo.alebo_acqf_optimizer(acq_function: AcquisitionFunction, bounds: Tensor, n: int, inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]], fixed_features: Optional[Dict[int, float]], rounding_func: Optional[Callable[[Tensor], Tensor]], raw_samples: int, num_restarts: int, B: Tensor) → Tuple[Tensor, 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[ALEBOGP, ModelListGP], objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]], X_observed: Tensor, X_pending: Optional[Tensor], q: int, noiseless: bool) → 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[ALEBOGP, ModelListGP], num_outputs: int) → List[MutableMapping[str, Tensor]][source]¶

Extract MAP statedict from the batch-mode ALEBO GP.

The batch GP can be either a single ALEBO GP or a ModelListGP of ALEBO GPs.

Parameters:

m_b – Batch-mode GP.
num_outputs – Number of outputs being modeled.

ax.models.torch.alebo.get_batch_model(B: Tensor, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor, Uvec_batch: Tensor, mean_constant_batch: Tensor, output_scale_batch: Tensor) → ALEBOGP[source]¶

Construct a batch-mode 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: Batch-mode ALEBO GP.

ax.models.torch.alebo.get_fitted_model(B: Tensor, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor, restarts: int, nsamp: int, init_state_dict: Optional[Dict[str, Tensor]]) → 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 batch-mode 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: Batch-mode (nsamp batches) fitted ALEBO GP.

ax.models.torch.alebo.get_map_model(B: Tensor, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor, restarts: int, init_state_dict: Optional[Dict[str, Tensor]]) → ExactMarginalLogLikelihood[source]¶

Do random-restart 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: non-batch ALEBO GP with MAP kernel hyperparameters.

ax.models.torch.alebo.laplace_sample_U(mll: ExactMarginalLogLikelihood, nsamp: int) → Tuple[Tensor, Tensor, 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 nsamp-1 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: ~typing.Callable[[~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[int], ~typing.List[int], ~typing.List[str], ~typing.Optional[~typing.Dict[str, ~torch.Tensor]], ~typing.Any], ~botorch.models.model.Model] = <function get_and_fit_model>, model_predictor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function predict_from_model>, acqf_constructor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Any], ~botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEI>, acqf_optimizer: ~typing.Callable[[~botorch.acquisition.acquisition.AcquisitionFunction, ~torch.Tensor, int, ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]], ~typing.Any], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: ~typing.Callable[[~ax.models.torch_base.TorchModel, ~typing.List[~typing.Tuple[float, float]], ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]]], ~typing.Optional[~typing.Dict[int, float]]], ~typing.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: ~typing.Any)[source]¶

Bases: 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 cross-validation.
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 model-specific options, and target_fidelities is a map from fidelity feature column indices to their respective target fidelities, used for multi-fidelity optimization problems. % TODO: refer to an example.

Xs: List[Tensor]¶

Ys: List[Tensor]¶

Yvars: List[Tensor]¶

best_point(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → Optional[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:

search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

d-tensor of the best point.

cross_validate(datasets: List[SupervisedDataset], X_test: Tensor, **kwargs: Any) → Tuple[Tensor, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
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.

Returns:

2-element 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[m1@j, m2@j].

feature_importances() → ndarray[source]¶

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

gen(n: int, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[source]¶

Generate new candidates.

Parameters:

n – Number of candidates to generate.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

A TorchGenResult container.

model: Optional[Model]¶

predict(X: Tensor) → Tuple[Tensor, Tensor][source]¶

Predict

Parameters:

X – (j x d) tensor of the j points at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

update(datasets: List[SupervisedDataset], 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome). None means that there is no additional data for the corresponding outcome.
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

ax.models.torch.botorch.get_rounding_func(rounding_func: Optional[Callable[[Tensor], Tensor]]) → Optional[Callable[[Tensor], Tensor]][source]¶

ax.models.torch.botorch_defaults module¶

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

ax.models.torch.botorch_defaults.get_and_fit_model(Xs: List[Tensor], Ys: List[Tensor], Yvars: List[Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, Tensor]] = None, refit_model: bool = True, use_input_warping: bool = False, use_loocv_pseudo_likelihood: bool = False, **kwargs: Any) → 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 if-else statements in lines 96-137. 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[Size] = None, task_feature: Optional[int] = None) → 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: TorchModel, bounds: List[Tuple[float, float]], objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, linear_constraints: Optional[Tuple[Tensor, Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[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 d-dimensional 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 model-specific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multi-fidelity optimization.

Returns:

A d-array of the best point, or None if no feasible point was observed.

ax.models.torch.botorch_defaults.recommend_best_out_of_sample_point(model: TorchModel, bounds: List[Tuple[float, float]], objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, linear_constraints: Optional[Tuple[Tensor, Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None, target_fidelities: Optional[Dict[int, float]] = None) → Optional[Tensor][source]¶

Identify the current best point by optimizing the posterior mean of the model. This is “out-of-sample” because it considers un-observed 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 d-dimensional 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 model-specific options.
target_fidelities – A map {feature_index: value} of fidelity feature column indices to their respective target fidelities. Used for multi-fidelity optimization.

Returns:

A d-array of the best point, or None if no feasible point exists.

ax.models.torch.botorch_defaults.scipy_optimizer(acq_function: AcquisitionFunction, bounds: Tensor, n: int, inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[Tensor], Tensor]] = None, **kwargs: Any) → Tuple[Tensor, Tensor][source]¶

Optimizer using scipy’s minimize module on a numpy-adpator.

Parameters:

acq_function – A botorch AcquisitionFunction.
bounds – A 2 x d-dim 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 round-trip transformations).

Returns:

2-element tuple containing

A n x d-dim 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 n-dim tensor of conditional acquisition values, where i-th element is the expected acquisition value conditional on having observed candidates 0,1,…,i-1.

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: 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 multi-fidelity 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 multi-fidelity optimization).
kwargs – Model-specific kwargs.

Xs: List[Tensor]¶

Ys: List[Tensor]¶

Yvars: List[Tensor]¶

fidelity_features: List[int]¶

gen(n: int, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[source]¶

Generate new candidates.

Parameters:

n – Number of candidates to generate.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

A TorchGenResults container, containing

(n x d) tensor of generated points.
n-tensor of weights for each point.
Dictionary of model-specific metadata for the given
generation candidates.

metric_names: List[str]¶

model: Optional[Model]¶

task_features: List[int]¶

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: BotorchModel

Max-value 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 multi-fidelity 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 multi-fidelity optimization).
kwargs – Model-specific kwargs.

Xs: List[Tensor]¶

Ys: List[Tensor]¶

Yvars: List[Tensor]¶

fidelity_features: List[int]¶

gen(n: int, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[source]¶

Generate new candidates.

Parameters:

n – Number of candidates to generate.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

A TorchGenResult container.

metric_names: List[str]¶

model: Optional[Model]¶

task_features: List[int]¶

ax.models.torch.botorch_moo module¶

class ax.models.torch.botorch_moo.MultiObjectiveBotorchModel(model_constructor: ~typing.Callable[[~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[int], ~typing.List[int], ~typing.List[str], ~typing.Optional[~typing.Dict[str, ~torch.Tensor]], ~typing.Any], ~botorch.models.model.Model] = <function get_and_fit_model>, model_predictor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function predict_from_model>, acqf_constructor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Any], ~botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEHVI>, acqf_optimizer: ~typing.Callable[[~botorch.acquisition.acquisition.AcquisitionFunction, ~torch.Tensor, int, ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]], ~typing.Any], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: ~typing.Callable[[~ax.models.torch_base.TorchModel, ~typing.List[~typing.Tuple[float, float]], ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]]], ~typing.Optional[~typing.Dict[int, float]]], ~typing.Optional[~torch.Tensor]] = <function recommend_best_observed_point>, frontier_evaluator: ~typing.Callable[[~ax.models.torch_base.TorchModel, ~torch.Tensor, ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]]], ~typing.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: ~typing.Any)[source]¶

Bases: BotorchModel

Customizable multi-objective 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[Tensor]¶

Ys: List[Tensor]¶

Yvars: List[Tensor]¶

gen(n: int, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[source]¶

Generate new candidates.

Parameters:

n – Number of candidates to generate.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

A TorchGenResult container.

ax.models.torch.botorch_moo_defaults module¶

References

[Daulton2020qehvi]

S. Daulton, M. Balandat, and E. Bakshy. Differentiable Expected Hypervolume Improvement for Parallel Multi-Objective 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: Model, objective_weights: Tensor, objective_thresholds: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None, X_pending: Optional[Tensor] = None, **kwargs: Any) → 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: Model, objective_weights: Tensor, objective_thresholds: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None, X_pending: Optional[Tensor] = None, **kwargs: Any) → 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[[TorchModel, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tuple[Tensor, Tensor]]], Tuple[Tensor, Tensor, Tensor]][source]¶

ax.models.torch.botorch_moo_defaults.get_weighted_mc_objective_and_objective_thresholds(objective_weights: Tensor, objective_thresholds: Tensor) → Tuple[WeightedMCMultiOutputObjective, 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 two-element tuple with the objective and objective thresholds

ax.models.torch.botorch_moo_defaults.infer_objective_thresholds(model: Model, objective_weights: Tensor, bounds: Optional[List[Tuple[float, float]]] = None, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, linear_constraints: Optional[Tuple[Tensor, Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, subset_idcs: Optional[Tensor] = None, Xs: Optional[List[Tensor]] = None, X_observed: Optional[Tensor] = None) → Tensor[source]¶

Infer objective thresholds.

This method uses the model-estimated Pareto frontier over the in-sample 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 in-sample 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 d-dimensional 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 d-dim tensor of in-sample points to use for determining the current in-sample Pareto frontier.

Returns:

A m-dim 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: TorchModel, objective_weights: Tensor, objective_thresholds: Optional[Tensor] = None, X: Optional[Tensor] = None, Y: Optional[Tensor] = None, Yvar: Optional[Tensor] = None, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None) → Tuple[Tensor, Tensor, 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:

3-element 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[m1@j, m2@j].
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[AcquisitionFunction], bounds: Tensor, inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[Tensor], Tensor]] = None, **kwargs: Any) → Tuple[Tensor, Tensor][source]¶

Sequential optimizer using scipy’s minimize module on a numpy-adaptor.

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 d-dim 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 round-trip transformations).

Returns:

2-element tuple containing

A n x d-dim tensor of generated candidates.
A n-dim tensor of conditional acquisition values, where i-th element is the expected acquisition value conditional on having observed candidates 0,1,…,i-1.

ax.models.torch.botorch_modular.acquisition module¶

class ax.models.torch.botorch_modular.acquisition.Acquisition(surrogate: Surrogate, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, botorch_acqf_class: Type[AcquisitionFunction], options: Optional[Dict[str, Any]] = None)[source]¶

Bases: 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).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).
botorch_acqf_class – Type of BoTorch AcquistitionFunction that should be used. Subclasses of Acquisition often specify these via default_botorch_acqf_class attribute, in which case specifying one here is not required.
options – Optional mapping of kwargs to the underlying Acquisition Function in BoTorch.

acqf: AcquisitionFunction¶

property botorch_acqf_class: Type[AcquisitionFunction]¶: BoTorch AcquisitionFunction class underlying this Acquisition.

compute_model_dependencies(surrogate: Surrogate, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, 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 non-empty 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).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).
options – The options kwarg dict, passed on initialization of the Acquisition object.

Returns: A dictionary of surrogate model-dependent options, to be passed: as kwargs to BoTorch`AcquisitionFunction` constructor.

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

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

evaluate(X: Tensor) → Tensor[source]¶

Evaluate the acquisition function on the candidate set X.

Parameters:: X – A batch_shape x q x d-dim Tensor of t-batches with q d-dim 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_and_transform(botorch_acqf_class: Type[AcquisitionFunction], model: Model, objective_weights: Tensor, objective_thresholds: Optional[Tensor] = None, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None) → Tuple[Optional[MCAcquisitionObjective], Optional[PosteriorTransform]][source]¶

property objective_thresholds: Optional[Tensor]¶

The objective thresholds for all outcomes.

For non-objective outcomes, the objective thresholds are nans.

property objective_weights: Optional[Tensor]¶: The objective weights for all outcomes.

optimize(n: int, search_space_digest: SearchSpaceDigest, inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, fixed_features: Optional[Dict[int, float]] = None, rounding_func: Optional[Callable[[Tensor], Tensor]] = None, optimizer_options: Optional[Dict[str, Any]] = None) → Tuple[Tensor, Tensor][source]¶

Generate a set of candidates via multi-start 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.
inequality_constraints – 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 post-processes an optimization result appropriately (i.e., according to round-trip transformations).
optimizer_options – Options for the optimizer function, e.g. sequential or raw_samples.

surrogate: Surrogate¶

ax.models.torch.botorch_modular.default_options module¶

ax.models.torch.botorch_modular.list_surrogate module¶

class ax.models.torch.botorch_modular.list_surrogate.ListSurrogate(botorch_submodel_class_per_outcome: ~typing.Optional[~typing.Dict[str, ~typing.Type[~botorch.models.model.Model]]] = None, botorch_submodel_class: ~typing.Optional[~typing.Type[~botorch.models.model.Model]] = None, submodel_options_per_outcome: ~typing.Optional[~typing.Dict[str, ~typing.Dict[str, ~typing.Any]]] = None, submodel_options: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None, mll_class: ~typing.Type[~gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood] = <class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>, mll_options: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None, submodel_outcome_transforms: ~typing.Optional[~typing.Dict[str, ~botorch.models.transforms.outcome.OutcomeTransform]] = None, submodel_input_transforms: ~typing.Optional[~typing.Dict[str, ~botorch.models.transforms.input.InputTransform]] = None, submodel_covar_module_class: ~typing.Optional[~typing.Dict[str, ~typing.Type[~gpytorch.kernels.kernel.Kernel]]] = None, submodel_covar_module_options: ~typing.Optional[~typing.Dict[str, ~typing.Dict[str, ~typing.Any]]] = None, submodel_likelihood_class: ~typing.Optional[~typing.Dict[str, ~typing.Type[~gpytorch.likelihoods.likelihood.Likelihood]]] = None, submodel_likelihood_options: ~typing.Optional[~typing.Dict[str, ~typing.Dict[str, ~typing.Any]]] = None)[source]¶

Bases: Surrogate

Special type of Surrogate that wraps a set of submodels into ModelListGP under the hood for multi-outcome or multi-task 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 underlying ModelListGP are of the same type. Use instead of botorch_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 model-fitting.
mll_options – Dictionary of options / kwargs for the MLL.
submodel_outcome_transforms – A dictionary mapping each outcome to a BoTorch outcome transform. Gets passed down to the BoTorch Model``s. To use multiple outcome transforms on a submodel, chain them together using ``ChainedOutcomeTransform.
submodel_input_transforms – A dictionary mapping each outcome to a BoTorch input transform. Gets passed down to the BoTorch Model. If sharing a single InputTransform object across submodels is preferred, pass in a dictionary where each outcome key references the same InputTransform object. To use multiple input transfroms on a submodel, chain them together using ChainedInputTransform.

botorch_submodel_class: Optional[Type[Model]]¶

botorch_submodel_class_per_outcome: Dict[str, Type[Model]]¶

construct(datasets: List[SupervisedDataset], metric_names: List[str], **kwargs: Any) → None[source]¶

Constructs the underlying BoTorch Model using the training data.

Parameters:

datasets – List of SupervisedDataset for the submodels of ModelListGP. Each training data is for one outcome, and the order of outcomes should match the order of metrics in metric_names argument.
metric_names – Names of metrics, in the same order as datasets (so if datasets is [ds_A, ds_B], the metrics are ["A" and "B"]). These are used to match training data with correct submodels of ModelListGP.
**kwargs –
Keyword arguments, accepts: - fidelity_features: Indices of columns in X that represent

fidelity
- task_features: Indices of columns in X that represent tasks

mll_class: Type[MarginalLogLikelihood]¶

mll_options: Dict[str, Any]¶

submodel_covar_module_class: Dict[str, Type[Kernel]]¶

submodel_covar_module_options: Dict[str, Dict[str, Any]]¶

submodel_input_transforms: Dict[str, InputTransform]¶

submodel_likelihood_class: Dict[str, Type[Likelihood]]¶

submodel_likelihood_options: Dict[str, Dict[str, Any]]¶

submodel_options: Dict[str, Any]¶

submodel_options_per_outcome: Dict[str, Dict[str, Any]]¶

submodel_outcome_transforms: Dict[str, OutcomeTransform]¶

ax.models.torch.randomforest module¶

class ax.models.torch.randomforest.RandomForest(max_features: Optional[str] = 'sqrt', num_trees: int = 500)[source]¶

Bases: TorchModel

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 one-hot 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(datasets: List[SupervisedDataset], X_test: Tensor, **kwargs: Any) → Tuple[Tensor, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
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.

Returns:

2-element 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[m1@j, m2@j].

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

predict(X: Tensor) → Tuple[Tensor, Tensor][source]¶

Predict

Parameters:

X – (j x d) tensor of the j points at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

ax.models.torch.botorch_modular.model module¶

class ax.models.torch.botorch_modular.model.BoTorchModel(acquisition_class: Optional[Type[Acquisition]] = None, acquisition_options: Optional[Dict[str, Any]] = None, botorch_acqf_class: Optional[Type[AcquisitionFunction]] = None, surrogate: Optional[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: TorchModel, 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, auto-selected 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, auto-selected 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 auto-selected based on experiment and data, with kwargs in surrogate_options applied.
surrogate_options – Optional dict of kwargs for Surrogate (used if no pre-instantiated 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: List[Tensor]¶

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 i-th model.

NOTE: This is an accessor for self.surrogate.Xs and returns it unchanged.

acquisition_class: Type[Acquisition]¶

acquisition_options: Dict[str, Any]¶

best_point(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → Optional[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:

search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

d-tensor of the best point.

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

cross_validate(datasets: List[SupervisedDataset], metric_names: List[str], X_test: Tensor, search_space_digest: SearchSpaceDigest) → Tuple[Tensor, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
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.

Returns:

2-element 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[m1@j, m2@j].

evaluate_acquisition_function(X: Tensor, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, acq_options: Optional[Dict[str, Any]] = None) → 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.
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).
acq_options – Keyword arguments used to contruct the acquisition function.

Returns:

A single-element tensor with the acquisition value for these points.

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, state_dict: Optional[Dict[str, Tensor]] = None, refit: bool = True) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

gen(n: int, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[source]¶

Generate new candidates.

Parameters:

n – Number of candidates to generate.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

A TorchGenResult container.

predict(X: Tensor) → Tuple[Tensor, Tensor][source]¶

Predict

Parameters:

X – (j x d) tensor of the j points at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

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

surrogate_options: Dict[str, Any]¶

update(datasets: List[Optional[SupervisedDataset]], metric_names: List[str], search_space_digest: SearchSpaceDigest, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome). None means that there is no additional data for the corresponding outcome.
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
candidate_metadata – Model-produced 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: Surrogate, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, botorch_acqf_class: Type[AcquisitionFunction], options: Optional[Dict[str, Any]] = None)[source]¶

Bases: Acquisition

acqf: AcquisitionFunction¶

compute_model_dependencies(surrogate: Surrogate, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, 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 non-empty 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).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).
options – The options kwarg dict, passed on initialization of the Acquisition object.

Returns: A dictionary of surrogate model-dependent options, to be passed: as kwargs to BoTorch`AcquisitionFunction` constructor.

surrogate: Surrogate¶

ax.models.torch.botorch_modular.optimizer_argparse module¶

ax.models.torch.botorch_modular.surrogate module¶

class ax.models.torch.botorch_modular.surrogate.Surrogate(botorch_model_class: ~typing.Type[~botorch.models.model.Model], model_options: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None, mll_class: ~typing.Type[~gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood] = <class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>, mll_options: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None, outcome_transform: ~typing.Optional[~botorch.models.transforms.outcome.OutcomeTransform] = None, input_transform: ~typing.Optional[~botorch.models.transforms.input.InputTransform] = None, covar_module_class: ~typing.Optional[~typing.Type[~gpytorch.kernels.kernel.Kernel]] = None, covar_module_options: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None, likelihood_class: ~typing.Optional[~typing.Type[~gpytorch.likelihoods.likelihood.Likelihood]] = None, likelihood_options: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None)[source]¶

Bases: 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 of BoTorchModel and is not meant to be used outside of it.

Parameters:

botorch_model_class – Model class to be used as the underlying BoTorch model.
model_options – Dictionary of options / kwargs for the BoTorch Model constructed during Surrogate.fit.
mll_class – MarginalLogLikelihood class to use for model-fitting.
mll_options – Dictionary of options / kwargs for the MLL.
outcome_transform – BoTorch outcome transforms. Passed down to the BoTorch Model. Multiple outcome transforms can be chained together using ChainedOutcomeTransform.
input_transform – BoTorch input transforms. Passed down to the BoTorch Model. Multiple input transforms can be chained together using ChainedInputTransform.
covar_module_class – Covariance module class, not yet used. Will be used to construct custom BoTorch Model in the future.
covar_module_options – Covariance module 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 BoTorch Model in the future.
likelihood_options – Likelihood options, not yet used. Will be used to construct custom BoTorch Model in the future.

property Xs: List[Tensor]¶

best_in_sample_point(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None) → Tuple[Tensor, float][source]¶: Finds the best observed point and the corresponding observed outcome values.

best_out_of_sample_point(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None) → Tuple[Tensor, Tensor][source]¶: Finds the best predicted point and the corresponding value of the appropriate best point acquisition function.

botorch_model_class: Type[Model]¶

clone_reset() → Surrogate[source]¶

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

construct(datasets: List[SupervisedDataset], **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 multi-output 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.

covar_module_class: Optional[Type[Kernel]] = None¶

covar_module_options: Dict[str, Any]¶

property device: device¶

property dtype: dtype¶

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, state_dict: Optional[Dict[str, Tensor]] = None, refit: bool = True) → None[source]¶

Fits the underlying BoTorch Model to m outcomes.

NOTE: state_dict and refit keyword arguments control how the undelying BoTorch Model will be fit: whether its parameters will be reoptimized and whether it will be warm-started from a given state.

There are three possibilities:

fit(state_dict=None): fit model from scratch (optimize model parameters and set its training data used for inference),
fit(state_dict=some_state_dict, refit=True): warm-start refit with a state dict of parameters (still re-optimize 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 cross-validation, for example).

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome), to be passed to Model.construct_inputs in BoTorch.
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.
state_dict – Optional state dict to load.
refit – Whether to re-optimize model parameters.

classmethod from_botorch(model: ~botorch.models.model.Model, mll_class: ~typing.Type[~gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood] = <class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>) → Surrogate[source]¶: Instantiate a Surrogate from a pre-instantiated Botorch Model.

input_transform: Optional[InputTransform] = None¶

likelihood_class: Optional[Type[Likelihood]] = None¶

likelihood_options: Dict[str, Any]¶

mll_class: Type[MarginalLogLikelihood]¶

mll_options: Dict[str, Any]¶

property model: Model¶

model_options: Dict[str, Any]¶

outcome_transform: Optional[OutcomeTransform] = None¶

pareto_frontier() → Tuple[Tensor, Tensor][source]¶

For multi-objective optimization, retrieve Pareto frontier instead of best point.

Returns: A two-tuple of:

tensor of points in the feature space,
tensor of corresponding (multiple) outcomes.

predict(X: Tensor) → Tuple[Tensor, 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 a n x o x o-dim tensor.

Return type:

Tensor

property training_data: List[SupervisedDataset]¶

update(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None, state_dict: Optional[Dict[str, Tensor]] = None, refit: bool = True) → None[source]¶

Updates the surrogate model with new data. In the base Surrogate, just calls fit after checking that this surrogate was not created via Surrogate.from_botorch (in which case the Model comes premade, constructed manually and then supplied to Surrogate).

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 – Model-produced metadata for candidates, in the order corresponding to the Xs.
state_dict – Optional state dict to load.
refit – Whether to re-optimize 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[Tensor]] = None, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, linear_constraints: Optional[Tuple[Tensor, Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, objective_thresholds: Optional[Tensor] = None, objective_weights: Optional[Tensor] = None) → Type[AcquisitionFunction][source]¶: Chooses a BoTorch AcquisitionFunction class.

ax.models.torch.botorch_modular.utils.choose_model_class(datasets: List[SupervisedDataset], search_space_digest: SearchSpaceDigest) → Type[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, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]], model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None) → Tuple[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]], Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]][source]¶: Extract acquisition and optimizer options from model_gen_options.

ax.models.torch.botorch_modular.utils.convert_to_block_design(datasets: List[SupervisedDataset], metric_names: List[str], force: bool = False) → Tuple[List[SupervisedDataset], List[str]][source]¶

ax.models.torch.botorch_modular.utils.fit_botorch_model(model: Union[Model, ModelList, ModelListGP], mll_class: Type[MarginalLogLikelihood], mll_options: Optional[Dict[str, Any]] = None) → None[source]¶: Fit a BoTorch model.

ax.models.torch.botorch_modular.utils.use_model_list(datasets: List[SupervisedDataset], botorch_model_class: Type[Model]) → bool[source]¶

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: BotorchModel

Does Bayesian optimization with Latent Context Embedding Additive (LCE-A) 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 pre-trained embeddings from embds_feature_dict only. Default is True.

Xs: List[Tensor]¶

Ys: List[Tensor]¶

Yvars: List[Tensor]¶

best_point(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → Optional[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:

search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

d-tensor of the best point.

fidelity_features: List[int]¶

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

get_and_fit_model(Xs: List[Tensor], Ys: List[Tensor], Yvars: List[Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, Tensor]] = None, fidelity_model_id: Optional[int] = None, **kwargs: Any) → 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.

metric_names: List[str]¶

model: Optional[Model]¶

task_features: List[int]¶

ax.models.torch.cbo_lcea.get_map_model(train_X: Tensor, train_Y: Tensor, train_Yvar: 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[LCEAGP, ExactMarginalLogLikelihood][source]¶: Obtain MAP fitting of Latent Context Embedding Additive (LCE-A) GP.

ax.models.torch.cbo_lcem module¶

class ax.models.torch.cbo_lcem.LCEMBO(context_cat_feature: Optional[Tensor] = None, context_emb_feature: Optional[Tensor] = None, embs_dim_list: Optional[List[int]] = None)[source]¶

Bases: BotorchModel

Does Bayesian optimization with LCE-M GP.

Xs: List[Tensor]¶

Ys: List[Tensor]¶

Yvars: List[Tensor]¶

fidelity_features: List[int]¶

get_and_fit_model(Xs: List[Tensor], Ys: List[Tensor], Yvars: List[Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, Tensor]] = None, fidelity_model_id: Optional[int] = None, **kwargs: Any) → ModelListGP[source]¶

Get a fitted multi-task 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.

metric_names: List[str]¶

model: Optional[Model]¶

task_features: List[int]¶

ax.models.torch.cbo_sac module¶

class ax.models.torch.cbo_sac.SACBO(decomposition: Dict[str, List[str]])[source]¶

Bases: 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]¶

fidelity_features: List[int]¶

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

get_and_fit_model(Xs: List[Tensor], Ys: List[Tensor], Yvars: List[Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, Tensor]] = None, fidelity_model_id: Optional[int] = None, **kwargs: Any) → 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.

metric_names: List[str]¶

model: Optional[Model]¶

task_features: List[int]¶

ax.models.torch.cbo_sac.generate_model_space_decomposition(decomposition: Dict[str, List[str]], feature_names: List[str]) → Dict[str, List[int]][source]¶

ax.models.torch.frontier_utils module¶

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

ax.models.torch.frontier_utils.get_weighted_mc_objective_and_objective_thresholds(objective_weights: Tensor, objective_thresholds: Tensor) → Tuple[WeightedMCMultiOutputObjective, 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 two-element 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. High-Dimensional Bayesian Optimization with Sparse Axis-Aligned Subspaces. Proceedings of the Thirty- Seventh Conference on Uncertainty in Artificial Intelligence, 2021.

[Eriksson2021nas]

D. Eriksson, P. Chuang, S. Daulton, et al. Latency-Aware Neural Architecture Search with Multi-Objective Bayesian Optimization. ICML AutoML Workshop, 2021.

class ax.models.torch.fully_bayesian.FullyBayesianBotorchModel(model_constructor: ~typing.Callable[[~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[int], ~typing.List[int], ~typing.List[str], ~typing.Optional[~typing.Dict[str, ~torch.Tensor]], ~typing.Any], ~botorch.models.model.Model] = <function get_and_fit_model_mcmc>, model_predictor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function predict_from_model_mcmc>, acqf_constructor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Any], ~botorch.acquisition.acquisition.AcquisitionFunction] = <function get_fully_bayesian_acqf>, acqf_optimizer: ~typing.Callable[[~botorch.acquisition.acquisition.AcquisitionFunction, ~torch.Tensor, int, ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]], ~typing.Any], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: ~typing.Callable[[~ax.models.torch_base.TorchModel, ~typing.List[~typing.Tuple[float, float]], ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]]], ~typing.Optional[~typing.Dict[int, float]]], ~typing.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: ~typing.Optional[bool] = None, num_samples: int = 256, warmup_steps: int = 512, thinning: int = 16, max_tree_depth: int = 6, disable_progbar: bool = False, gp_kernel: str = 'matern', verbose: bool = False, **kwargs: ~typing.Any)[source]¶

Bases: FullyBayesianBotorchModelMixin, BotorchModel

Fully Bayesian Model that uses NUTS to sample from hyperparameter posterior.

This includes support for using sparse axis-aligned subspace priors (SAAS). See [Eriksson2021saasbo] for details.

class ax.models.torch.fully_bayesian.FullyBayesianBotorchModelMixin[source]¶

Bases: object

feature_importances() → ndarray[source]¶

model: Optional[Model] = None¶

class ax.models.torch.fully_bayesian.FullyBayesianMOOBotorchModel(model_constructor: ~typing.Callable[[~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[int], ~typing.List[int], ~typing.List[str], ~typing.Optional[~typing.Dict[str, ~torch.Tensor]], ~typing.Any], ~botorch.models.model.Model] = <function get_and_fit_model_mcmc>, model_predictor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function predict_from_model_mcmc>, acqf_constructor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Any], ~botorch.acquisition.acquisition.AcquisitionFunction] = <function get_fully_bayesian_acqf_nehvi>, acqf_optimizer: ~typing.Callable[[~botorch.acquisition.acquisition.AcquisitionFunction, ~torch.Tensor, int, ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]], ~typing.Any], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: ~typing.Callable[[~ax.models.torch_base.TorchModel, ~typing.List[~typing.Tuple[float, float]], ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]]], ~typing.Optional[~typing.Dict[int, float]]], ~typing.Optional[~torch.Tensor]] = <function recommend_best_observed_point>, frontier_evaluator: ~typing.Callable[[~ax.models.torch_base.TorchModel, ~torch.Tensor, ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]]], ~typing.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 = 256, warmup_steps: int = 512, thinning: int = 16, max_tree_depth: int = 6, use_saas: ~typing.Optional[bool] = None, disable_progbar: bool = False, gp_kernel: str = 'matern', verbose: bool = False, **kwargs: ~typing.Any)[source]¶

Bases: FullyBayesianBotorchModelMixin, 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: Tensor, Z: Tensor, lengthscale: Tensor) → 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[Tensor], Ys: List[Tensor], Yvars: List[Tensor], task_features: List[int], fidelity_features: List[int], metric_names: List[str], state_dict: Optional[Dict[str, Tensor]] = None, refit_model: bool = True, use_input_warping: bool = False, use_loocv_pseudo_likelihood: bool = False, num_samples: int = 256, warmup_steps: int = 512, thinning: int = 16, max_tree_depth: int = 6, disable_progbar: bool = False, gp_kernel: str = 'matern', verbose: bool = False, **kwargs: Any) → 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: ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = None, X_observed: ~typing.Optional[~torch.Tensor] = None, X_pending: ~typing.Optional[~torch.Tensor] = None, acqf_constructor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Any], ~botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEI>, **kwargs: ~typing.Any) → AcquisitionFunction[source]¶

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

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

ax.models.torch.fully_bayesian.predict_from_model_mcmc(model: Model, X: Tensor) → Tuple[Tensor, 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 o-dim tensor. Tensor: The predicted posterior covariance as a n x o x o-dim tensor.

Return type:

Tensor

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

ax.models.torch.fully_bayesian.run_inference(pyro_model: Callable, X: Tensor, Y: Tensor, Yvar: Tensor, num_samples: int = 256, warmup_steps: int = 512, 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, Tensor][source]¶

ax.models.torch.fully_bayesian.single_task_pyro_model(X: Tensor, Y: Tensor, Yvar: Tensor, use_input_warping: bool = False, eps: float = 1e-07, 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: GPyTorchModel, mcmc_samples: Dict) → None[source]¶: Load MCMC samples into GPyTorchModel.

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

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

ax.models.torch.fully_bayesian_model_utils.pyro_sample_noise(**tkwargs: Any) → Tensor[source]¶

ax.models.torch.fully_bayesian_model_utils.pyro_sample_outputscale(concentration: float = 2.0, rate: float = 0.15, **tkwargs: Any) → Tensor[source]¶

ax.models.torch.fully_bayesian_model_utils.pyro_sample_saas_lengthscales(dim: int, alpha: float = 0.1, **tkwargs: Any) → Tensor[source]¶

ax.models.torch.posterior_mean module¶

ax.models.torch.posterior_mean.get_PosteriorMean(model: Model, objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None, X_pending: Optional[Tensor] = None, **kwargs: Any) → 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: Tensor, initial_X_d: Tensor, bounds_d: List[Tuple[float, float]], **kwargs: Any)[source]¶

Bases: 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 low-d space must also be provided, which in the REMBO paper should be [(-sqrt(d), sqrt(d)]^d.

Function evaluations happen in the high-D space, and so the arms on the experiment will also be tracked in the high-D space. This class maintains a list of points in the low-d spac that have been launched, so we can match arms in high-D space back to their low-d point on update.

Parameters:

A – (D x d) projection matrix.
initial_X_d – Points in low-d space for initial data.
bounds_d – Box bounds in the low-d space.
kwargs – kwargs for BotorchModel init

Xs: List[Tensor]¶

Ys: List[Tensor]¶

Yvars: List[Tensor]¶

best_point(search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → Optional[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:

search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

d-tensor of the best point.

cross_validate(datasets: List[SupervisedDataset], X_test: Tensor, **kwargs: Any) → Tuple[Tensor, 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
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.

Returns:

2-element 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[m1@j, m2@j].

fidelity_features: List[int]¶

fit(datasets: List[SupervisedDataset], metric_names: List[str], search_space_digest: SearchSpaceDigest, candidate_metadata: Optional[List[List[Optional[Dict[str, Any]]]]] = None) → None[source]¶

Fit model to m outcomes.

Parameters:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in the datasets.
candidate_metadata – Model-produced metadata for candidates, in the order corresponding to the Xs.

from_01(X_d01: Tensor) → 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, search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig) → TorchGenResults[source]¶

Generate new candidates.

Parameters:

n – Number of candidates to generate.
search_space_digest – A SearchSpaceDigest object containing metadata about the search space (e.g. bounds, parameter types).
torch_opt_config – A TorchOptConfig object containing optimization arguments (e.g., objective weights, constraints).

Returns:

A TorchGenResult container.

metric_names: List[str]¶

model: Optional[Model]¶

predict(X: Tensor) → Tuple[Tensor, Tensor][source]¶

Predict

Parameters:

X – (j x d) tensor of the j points at which to make predictions.

Returns:

2-element 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[m1@j, m2@j].

project_down(X_D: Tensor) → Tensor[source]¶

Map points in the high-D space to the low-d 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 high-D space.
Returns:: Tensor in low-d space.
Return type:: X_d

project_up(X: Tensor) → Tensor[source]¶: Project to high-dimensional space.

task_features: List[int]¶

to_01(X_d: Tensor) → Tensor[source]¶

Map points from bounds_d to [0, 1].

Parameters:: X_d – Tensor in bounds_d

Returns: Tensor in [0, 1].

update(datasets: List[SupervisedDataset], 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:

datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome). None means that there is no additional data for the corresponding outcome.
metric_names – A list of metric names, with the i-th metric corresponding to the i-th dataset.
search_space_digest – A SearchSpaceDigest object containing metadata on the features in X.
candidate_metadata – Model-produced 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: Tensor¶

model: Model¶

objective_thresholds: Optional[Tensor]¶

objective_weights: Tensor¶

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

ax.models.torch.utils.get_botorch_objective_and_transform(model: Model, objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, objective_thresholds: Optional[Tensor] = None, X_observed: Optional[Tensor] = None) → Tuple[Optional[MCAcquisitionObjective], Optional[PosteriorTransform]][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.
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:

A two-tuple containing (optioally) an MCAcquisitionObjective and (optionally) a PosteriorTransform.

ax.models.torch.utils.get_out_of_sample_best_point_acqf(model: Model, Xs: List[Tensor], X_observed: Tensor, objective_weights: 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[Tensor, Tensor]] = None, seed_inner: Optional[int] = None, qmc: bool = True, **kwargs: Any) → Tuple[AcquisitionFunction, Optional[List[int]]][source]¶

Picks an appropriate acquisition function to find the best out-of-sample (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: Model) → bool[source]¶: Check if a given (single-task) 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 d-1.

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

ax.models.torch.utils.predict_from_model(model: Model, X: Tensor) → Tuple[Tensor, Tensor][source]¶

Predicts outcomes given a model and input tensor.

For a FullyBayesianPosterior we currently use a Gaussian approximation where we compute the mean and variance of the Gaussian mixture. This should ideally be changed to compute quantiles instead when Ax supports non-Gaussian distributions.

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 a n x o x o-dim tensor.

Return type:

Tensor

ax.models.torch.utils.randomize_objective_weights(objective_weights: Tensor, **acquisition_function_kwargs: Any) → 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 d-1.

ax.models.torch.utils.subset_model(model: Model, objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, objective_thresholds: Optional[Tensor] = None) → 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 null-op 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 m-dim 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.

ax.models.torch.utils.tensor_callable_to_array_callable(tensor_func: Callable[[Tensor], Tensor], device: device) → Callable[[ndarray], ndarray][source]¶: transfer a tensor callable to an array callable