ax.models¶

ax.models.base¶

class ax.models.base.Model[source]¶

Bases: object

Base class for an Ax model.

Note: the core methods each model has: fit, predict, gen, cross_validate, and best_point are not present in this base class, because the signatures for those methods vary based on the type of the model. This class only contains the methods that all models have in common and for which they all share the signature.

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[Union[NoneType, str, bool, float, int]]], objective_weights: Optional[ndarray], outcome_constraints: Optional[tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[dict[int, Union[NoneType, str, bool, float, int]]] = None, pending_observations: Optional[list[list[list[Union[NoneType, str, bool, float, int]]]]] = None, model_gen_options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = None) → Optional[list[Union[NoneType, 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[Union[NoneType, str, bool, float, int]]]], Ys_train: list[list[float]], Yvars_train: list[list[float]], X_test: list[list[Union[NoneType, str, bool, float, int]]], use_posterior_predictive: bool = False) → tuple[numpy.ndarray, numpy.ndarray][source]¶

Do cross validation with the given training and test sets.

Training set is given in the same format as to fit. Test set is given in the same format as to predict.

Parameters:

• Xs_train – A list of m lists X of parameterizations (each parameterization is a list of parameter values of length d), each of length k_i, for each outcome.

• Ys_train – The corresponding list of m lists Y, each of length k_i, for each outcome.

• Yvars_train – The variances of each entry in Ys, same shape.

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

• use_posterior_predictive – A boolean indicating if the predictions should be from the posterior predictive (i.e. including observation noise).

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[Union[NoneType, str, bool, float, int]]]], Ys: list[list[float]], Yvars: list[list[float]], parameter_values: list[list[Union[NoneType, 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[Union[NoneType, str, bool, float, int]]], objective_weights: Optional[ndarray], outcome_constraints: Optional[tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[dict[int, Union[NoneType, str, bool, float, int]]] = None, pending_observations: Optional[list[list[list[Union[NoneType, str, bool, float, int]]]]] = None, model_gen_options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = None) → tuple[list[list[Union[NoneType, 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[Union[NoneType, str, bool, float, int]]]) → tuple[numpy.ndarray, numpy.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: dict[str, typing.Any] = <factory>, candidate_metadata: ~typing.Optional[list[typing.Optional[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[botorch.utils.datasets.SupervisedDataset], X_test: Tensor, search_space_digest: SearchSpaceDigest, use_posterior_predictive: bool = False) → tuple[torch.Tensor, torch.Tensor][source]¶

Do cross validation with the given training and test sets.

Training set is given in the same format as to fit. Test set is given in the same format as to predict.

Parameters:

• datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).

• 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.

• use_posterior_predictive – A boolean indicating if the predictions should be from the posterior predictive (i.e. including observation noise).

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[botorch.utils.datasets.SupervisedDataset], 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).

• 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[torch.Tensor, torch.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[botorch.utils.datasets.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[tuple[torch.Tensor, torch.Tensor]] = None, objective_thresholds: ~typing.Optional[~torch.Tensor] = None, linear_constraints: ~typing.Optional[tuple[torch.Tensor, torch.Tensor]] = None, fixed_features: ~typing.Optional[dict[int, float]] = None, pending_observations: ~typing.Optional[list[torch.Tensor]] = None, model_gen_options: dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, typing.Any], dict[str, typing.Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]] = <factory>, rounding_func: ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]] = None, opt_config_metrics: ~typing.Optional[dict[str, ax.core.metric.Metric]] = None, is_moo: bool = False, risk_measure: ~typing.Optional[~botorch.acquisition.risk_measures.RiskMeasureMCObjective] = None, fit_out_of_design: 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. This commonly includes optimizer_kwargs, which often specifies the optimizer options to be passed to the optimizer while optimizing the acquisition function. These are generally expected to mimic the signature of optimize_acqf, though not all models may support all possible arguments and some models may support additional arguments that are not passed to the optimizer. While constructing a generation strategy, these options can be passed in as follows: >>> model_gen_kwargs = { >>> “model_gen_options”: { >>> “optimizer_kwargs”: { >>> “num_restarts”: 20, >>> “sequential”: False, >>> “options”: { >>> “batch_limit: 5, >>> “maxiter”: 200, >>> }, >>> }, >>> }, >>> }

Type:

dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]

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

risk_measure¶

An optional risk measure, used for robust optimization.

Type:

Optional[botorch.acquisition.risk_measures.RiskMeasureMCObjective]

fit_out_of_design: bool = False¶

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

is_moo: bool = False¶

linear_constraints: Optional[tuple[torch.Tensor, torch.Tensor]] = None¶

model_gen_options: dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]¶

objective_thresholds: Optional[Tensor] = None¶

objective_weights: Tensor¶

opt_config_metrics: Optional[dict[str, ax.core.metric.Metric]] = None¶

outcome_constraints: Optional[tuple[torch.Tensor, torch.Tensor]] = None¶

pending_observations: Optional[list[torch.Tensor]] = None¶

risk_measure: Optional[RiskMeasureMCObjective] = None¶

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

ax.models.model_utils module¶

class ax.models.model_utils.TorchModelLike(*args, **kwargs)[source]¶

Bases: Protocol

A protocol that stands in for TorchModel like objects that have a predict method.

predict(X: Tensor) → tuple[torch.Tensor, torch.Tensor][source]¶

Predicts outcomes given an input tensor.

Parameters:

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.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[torch.Tensor, numpy.ndarray], ...]]) → Union[ndarray, tuple[numpy.ndarray, ...]][source]¶

Convert every item in a tuple of tensors/arrays into an array.

Parameters:

x – A tensor, array, or a tuple of potentially mixed tensors and arrays.

Returns:

x, with everything converted to array.

ax.models.model_utils.best_in_sample_point(Xs: Union[list[torch.Tensor], list[numpy.ndarray]], model: TorchModelLike, bounds: list[tuple[float, float]], objective_weights: Optional[Union[Tensor, ndarray]], outcome_constraints: Optional[tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, linear_constraints: Optional[tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, fixed_features: Optional[dict[int, float]] = None, risk_measure: Optional[RiskMeasureMCObjective] = None, options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = None) → Optional[tuple[Union[torch.Tensor, numpy.ndarray], float]][source]¶

Select the best point that has been observed.

Implements two approaches to selecting the best point.

For both approaches, only points that satisfy parameter space constraints (bounds, linear_constraints, fixed_features) will be returned. Points must also be observed for all objective and constraint outcomes. Returned points may violate outcome constraints, depending on the method below.

1: Select the point that maximizes the expected utility (objective_weights^T posterior_objective_means - baseline) * Prob(feasible) Here baseline should be selected so that at least one point has positive utility. It can be specified in the options dict, otherwise min (objective_weights^T posterior_objective_means) will be used, where the min is over observed points.

2: Select the 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 – A Torch model or Surrogate.

• 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.

• risk_measure – An optional risk measure for reporting best robust 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: TorchModelLike, bounds: list[tuple[float, float]], objective_weights: Optional[Union[Tensor, ndarray]], outcome_constraints: Optional[tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, linear_constraints: Optional[tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]] = None, fixed_features: Optional[dict[int, float]] = None, risk_measure: Optional[RiskMeasureMCObjective] = None, options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = 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 – A Torch model or Surrogate.

• 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.

• risk_measure – An optional risk measure for reporting best robust 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[numpy.ndarray, numpy.ndarray], point: ndarray) → tuple[bool, numpy.ndarray][source]¶

Check if a point satisfies parameter constraints.

Parameters:

• linear_constraints – A tuple of (A, b). For k linear constraints on 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: Mapping[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[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.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[torch.Tensor], list[numpy.ndarray]], objective_weights: Union[Tensor, ndarray], outcome_constraints: Optional[tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.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) → Mapping[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[numpy.ndarray, numpy.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[numpy.ndarray, int][source]¶

Rejection sample in parameter space. Parameter space is typically [0, 1] for all tunable parameters.

Models must implement a gen_unconstrained method in order to support rejection sampling via this utility.

Parameters:

• gen_unconstrained – A callable that generates unconstrained points in the parameter space. This is typically the _gen_unconstrained method of a RandomModel.

• n – Number of samples to generate.

• d – Dimensionality of the parameter space.

• tunable_feature_indices – Indices of the tunable features in the parameter space.

• 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.

• deduplicate – If true, reject points that are duplicates of previously generated points. The points are deduplicated after applying the rounding function.

• max_draws – Maximum number of attemped draws before giving up.

• 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).

• existing_points – A set of previously generated points to use for deduplication. These should be provided in the parameter space model operates in.

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¶

ax.models.winsorization_config module¶

class ax.models.winsorization_config.WinsorizationConfig(lower_quantile_margin: float = 0.0, upper_quantile_margin: float = 0.0, lower_boundary: Optional[float] = None, upper_boundary: Optional[float] = None)[source]¶

Bases: object

Dataclass for storing Winsorization configuration parameters

Attributes: lower_quantile_margin: Winsorization will increase any metric value below this

quantile to this quantile’s value.

upper_quantile_margin: Winsorization will decrease any metric value above this

quantile to this quantile’s value. NOTE: this quantile will be inverted before any operations, e.g., a value of 0.2 will decrease values above the 80th percentile to the value of the 80th percentile.

lower_boundary: If this value is lesser than the metric value corresponding to

lower_quantile_margin, set metric values below lower_boundary to lower_boundary and leave larger values unaffected.

upper_boundary: If this value is greater than the metric value corresponding to

upper_quantile_margin, set metric values above upper_boundary to upper_boundary and leave smaller values unaffected.

lower_boundary: Optional[float] = None¶

lower_quantile_margin: float = 0.0¶

upper_boundary: Optional[float] = None¶

upper_quantile_margin: float = 0.0¶

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[Union[NoneType, str, bool, float, int]]], objective_weights: Optional[ndarray], outcome_constraints: Optional[tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[dict[int, Union[NoneType, str, bool, float, int]]] = None, pending_observations: Optional[list[list[list[Union[NoneType, str, bool, float, int]]]]] = None, model_gen_options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = None) → tuple[list[list[Union[NoneType, 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[Union[NoneType, str, bool, float, int]]]], Ys: list[list[float]], Yvars: list[list[float]], parameter_values: list[list[Union[NoneType, 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[Union[NoneType, str, bool, float, int]]], objective_weights: Optional[ndarray], outcome_constraints: Optional[tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[dict[int, Union[NoneType, str, bool, float, int]]] = None, pending_observations: Optional[list[list[list[Union[NoneType, str, bool, float, int]]]]] = None, model_gen_options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = None) → tuple[list[list[Union[NoneType, 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[Union[NoneType, str, bool, float, int]]]) → tuple[numpy.ndarray, numpy.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.random.base module¶

class ax.models.random.base.RandomModel(deduplicate: bool = True, seed: Optional[int] = None, init_position: int = 0, 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.

seed¶

An optional seed value for scrambling.

init_position¶

The initial state of the generator. This is the number of samples to fast-forward before generating new samples. Used to ensure that the re-loaded generator will continue generating from the same sequence rather than starting from scratch.

generated_points¶

A set of previously generated points to use for deduplication. These should be provided in the raw transformed space the model operates in.

fallback_to_sample_polytope¶

If True, when rejection sampling fails, we fall back to the HitAndRunPolytopeSampler.

gen(n: int, bounds: list[tuple[float, float]], linear_constraints: Optional[tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[dict[int, float]] = None, model_gen_options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = None, rounding_func: Optional[Callable[[ndarray], ndarray]] = None) → tuple[numpy.ndarray, numpy.ndarray][source]¶

Generate new candidates.

Parameters:

• n – Number of candidates to generate.

• bounds – A list of (lower, upper) tuples for each column of X. Defined on [0, 1]^d.

• linear_constraints – A tuple of (A, b). For k linear constraints on 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 = True, seed: Optional[int] = None, init_position: int = 0, generated_points: Optional[ndarray] = None, fallback_to_sample_polytope: bool = False)[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.

See base RandomModel for a description of model attributes.

ax.models.random.sobol module¶

class ax.models.random.sobol.SobolGenerator(deduplicate: bool = True, seed: Optional[int] = None, 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.

scramble¶

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

See base `RandomModel` for a description of remaining attributes.

property engine: Optional[SobolEngine]¶

Return a singleton SobolEngine.

gen(n: int, bounds: list[tuple[float, float]], linear_constraints: Optional[tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[dict[int, float]] = None, model_gen_options: Optional[dict[str, Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, Any], dict[str, Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]] = None, rounding_func: Optional[Callable[[ndarray], ndarray]] = None) → tuple[numpy.ndarray, numpy.ndarray][source]¶

Generate new candidates.

Parameters:

• n – Number of candidates to generate.

• bounds – A list of (lower, upper) tuples for each column of X.

• linear_constraints – A tuple of (A, b). For k linear constraints on 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).

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.torch.botorch module¶

class ax.models.torch.botorch.BotorchModel(model_constructor: ~typing.Callable[[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[int], list[int], list[str], ~typing.Optional[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, bool], tuple[torch.Tensor, torch.Tensor]] = <function predict_from_model>, acqf_constructor: ~ax.models.torch.botorch_defaults.TAcqfConstructor = <function get_qLogNEI>, acqf_optimizer: ~typing.Callable[[~botorch.acquisition.acquisition.AcquisitionFunction, ~torch.Tensor, int, ~typing.Optional[list[tuple[torch.Tensor, torch.Tensor, float]]], ~typing.Optional[list[tuple[torch.Tensor, torch.Tensor, float]]], ~typing.Optional[dict[int, float]], ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]], ~typing.Any], tuple[torch.Tensor, torch.Tensor]] = <function scipy_optimizer>, best_point_recommender: ~typing.Callable[[~ax.models.torch_base.TorchModel, list[tuple[float, float]], ~torch.Tensor, ~typing.Optional[tuple[torch.Tensor, torch.Tensor]], ~typing.Optional[tuple[torch.Tensor, torch.Tensor]], ~typing.Optional[dict[int, float]], ~typing.Optional[dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, list[str], dict[int, typing.Any], dict[str, typing.Any], ax.core.optimization_config.OptimizationConfig, ax.models.winsorization_config.WinsorizationConfig, NoneType]]], ~typing.Optional[dict[int, float]]], ~typing.Optional[~torch.Tensor]] = <function recommend_best_observed_point>, refit_on_cv: bool = False, warm_start_refitting: bool = True, use_input_warping: bool = False, use_loocv_pseudo_likelihood: bool = False, prior: ~typing.Optional[dict[str, typing.Any]] = None, **kwargs: ~typing.Any)[source]¶

Bases: TorchModel

Customizable botorch model.

By default, this uses a noisy Log Expected Improvement (qLogNEI) 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.

• warm_start_refitting – If True, start model refitting from previous model parameters in order to speed up the fitting process.

• prior –

  An optional dictionary that contains the specification of GP model prior. Currently, the keys include: - covar_module_prior: prior on covariance matrix e.g.

  {“lengthscale_prior”: GammaPrior(3.0, 6.0)}.

  • type: type of prior on task covariance matrix e.g.`LKJCovariancePrior`.

  • sd_prior: A scalar prior over nonnegative numbers, which is used for the

    default LKJCovariancePrior task_covar_prior.

  • eta: The eta parameter on the default LKJ task_covar_prior.

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_qLogNEI.

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[torch.Tensor]¶

Ys: list[torch.Tensor]¶

Yvars: list[torch.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[botorch.utils.datasets.SupervisedDataset], X_test: Tensor, use_posterior_predictive: bool = False, **kwargs: Any) → tuple[torch.Tensor, torch.Tensor][source]¶

Do cross validation with the given training and test sets.

Training set is given in the same format as to fit. Test set is given in the same format as to predict.

Parameters:

• datasets – A list of SupervisedDataset containers, each corresponding to the data of one metric (outcome).

• 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.

• use_posterior_predictive – A boolean indicating if the predictions should be from the posterior predictive (i.e. including observation noise).

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[botorch.utils.datasets.SupervisedDataset], 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).

• 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.

property model: Model¶

predict(X: Tensor) → tuple[torch.Tensor, torch.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 search_space_digest: SearchSpaceDigest¶