# ax.models¶

## Base Models¶

### ax.models.discrete_base module¶

class ax.models.discrete_base.DiscreteModel[source]

Bases: ax.models.base.Model

This class specifies the interface for a model based on discrete parameters.

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

best_point(n, parameter_values, objective_weights, outcome_constraints=None, fixed_features=None, pending_observations=None, model_gen_options=None)[source]

Obtains the point that has the best value according to the model prediction and its model predictions.

Return type

Optional[List[Union[str, float, int, None]]]

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, Ys_train, Yvars_train, X_test)[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
Return type

Tuple[ndarray, ndarray]

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, Ys, Yvars, parameter_values, outcome_names)[source]

Fit model to m outcomes.

Parameters
Return type

None

gen(n, parameter_values, objective_weights, outcome_constraints=None, fixed_features=None, pending_observations=None, model_gen_options=None)[source]

Generate new candidates.

Parameters
Return type
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)[source]

Predict

Parameters

X (List[List[Union[str, float, int, None]]]) – List of the j parameterizations at which to make predictions.

Return type

Tuple[ndarray, ndarray]

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

update(Xs, Ys, Yvars, parameter_values, outcome_names)[source]

Update model with new data. Note that this data will be treated as new even if the model has already been fit to it, so ensure that only new data is passed in or use fit to refit the model from scratch.

Parameters
Return type

None

### ax.models.model_utils module¶

ax.models.model_utils.add_fixed_features(tunable_points, d, fixed_features, tunable_feature_indices)[source]

Add fixed features to points in tunable space.

Parameters
Returns

Points in the full d-dimensional space, defined by bounds.

Return type

points

ax.models.model_utils.as_array(x)[source]

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

Parameters

x (Union[Tensor, ndarray, Tuple[Union[Tensor, ndarray], …]]) – A tensor, array, or a tuple of potentially mixed tensors and arrays.

Return type

Union[ndarray, Tuple[ndarray, …]]

Returns

x, with everything converted to array.

ax.models.model_utils.best_observed_point(model, bounds, objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=None, options=None)[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
Return type

Optional[ndarray]

Returns

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

ax.models.model_utils.check_duplicate(point, points)[source]

Check if a point exists in another array.

Parameters
• point (ndarray) – Newly generated point to check.

• prev_points – Points previously generated.

Return type

bool

Returns

True if the point is contained in points, else False

ax.models.model_utils.check_param_constraints(linear_constraints, point)[source]

Check if a point satisfies parameter constraints.

Parameters
• linear_constraints (Tuple[ndarray, ndarray]) – 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 (ndarray) – A candidate point in d-dimensional space, as a (1 x d) matrix.

Return type

Tuple[bool, ndarray]

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.filter_constraints_and_fixed_features(X, bounds, linear_constraints=None, fixed_features=None)[source]

Filter points to those that satisfy bounds, linear_constraints, and fixed_features.

Parameters
Return type

Union[Tensor, ndarray]

Returns

Feasible points.

ax.models.model_utils.get_observed(Xs, objective_weights, outcome_constraints=None)[source]

Filter points to those that are observed for objective outcomes and outcomes that show up in outcome_constraints (if there are any).

Parameters
Return type

Union[Tensor, ndarray]

Returns

Points observed for all objective outcomes and outcome constraints.

ax.models.model_utils.rejection_sample(gen_unconstrained, n, d, tunable_feature_indices, linear_constraints=None, deduplicate=False, max_draws=None, fixed_features=None, rounding_func=None, existing_points=None)[source]

Rejection sample in parameter space.

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

Return type

Tuple[ndarray, int]

ax.models.model_utils.tunable_feature_indices(bounds, fixed_features=None)[source]

Get the feature indices of tunable features.

Parameters
Return type

ndarray

Returns

The indices of tunable features.

ax.models.model_utils.validate_bounds(bounds, fixed_feature_indices)[source]

Ensure the requested space is [0,1]^d.

Parameters
Return type

None

### ax.models.numpy_base module¶

class ax.models.numpy_base.NumpyModel[source]

Bases: ax.models.base.Model

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

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

best_point(bounds, objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=None, model_gen_options=None)[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
Return type

Optional[ndarray]

Returns

A d-array of the best point.

cross_validate(Xs_train, Ys_train, Yvars_train, X_test)[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 (List[ndarray]) – A list of m (k_i x d) feature matrices X. Number of rows k_i can vary from i=1,…,m.

• Ys_train (List[ndarray]) – The corresponding list of m (k_i x 1) outcome arrays Y, for each outcome.

• Yvars_train (List[ndarray]) – The variances of each entry in Ys, same shape.

• X_test (ndarray) – (j x d) array of the j points at which to make predictions.

Return type

Tuple[ndarray, ndarray]

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, Ys, Yvars, bounds, task_features, feature_names, fidelity_features)[source]

Fit model to m outcomes.

Parameters
Return type

None

gen(n, bounds, objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=None, pending_observations=None, model_gen_options=None, rounding_func=None)[source]

Generate new candidates.

Parameters
Return type

Tuple[ndarray, ndarray]

Returns

2-element tuple containing

• (n x d) array of generated points.

• n-array of weights for each point.

predict(X)[source]

Predict

Parameters

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

Return type

Tuple[ndarray, ndarray]

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

update(Xs, Ys, Yvars)[source]

Update the model.

Updating the model requires both existing and additional data. The data passed into this method will become the new training data.

Parameters
• Xs (List[ndarray]) – Existing + additional data for the model, in the same format as for fit.

• Ys (List[ndarray]) – Existing + additional data for the model, in the same format as for fit.

• Yvars (List[ndarray]) – Existing + additional data for the model, in the same format as for fit.

Return type

None

### ax.models.torch_base module¶

class ax.models.torch_base.TorchModel[source]

Bases: ax.models.base.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(bounds, objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=None, model_gen_options=None)[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
Return type

Optional[Tensor]

Returns

d-tensor of the best point.

cross_validate(Xs_train, Ys_train, Yvars_train, X_test)[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 (List[Tensor]) – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.

• Ys_train (List[Tensor]) – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.

• Yvars_train (List[Tensor]) – The variances of each entry in Ys, same shape.

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

Return type

Tuple[Tensor, Tensor]

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(Xs, Ys, Yvars, bounds, task_features, feature_names, fidelity_features)[source]

Fit model to m outcomes.

Parameters
Return type

None

gen(n, bounds, objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=None, pending_observations=None, model_gen_options=None, rounding_func=None)[source]

Generate new candidates.

Parameters
Return type

Tuple[Tensor, Tensor]

Returns

2-element tuple containing

• (n x d) tensor of generated points.

• n-tensor of weights for each point.

predict(X)[source]

Predict

Parameters

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

Return type

Tuple[Tensor, Tensor]

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(Xs, Ys, Yvars)[source]

Update the model.

Updating the model requires both existing and additional data. The data passed into this method will become the new training data.

Parameters
• Xs (List[Tensor]) – Existing + additional data for the model, in the same format as for fit.

• Ys (List[Tensor]) – Existing + additional data for the model, in the same format as for fit.

• Yvars (List[Tensor]) – Existing + additional data for the model, in the same format as for fit.

Return type

None

## Discrete Models¶

### ax.models.discrete.eb_thompson module¶

class ax.models.discrete.eb_thompson.EmpiricalBayesThompsonSampler(num_samples=10000, min_weight=None, uniform_weights=False)[source]

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=100, check_cardinality=True)[source]

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, parameter_values, objective_weights, outcome_constraints=None, fixed_features=None, pending_observations=None, model_gen_options=None)[source]

Generate new candidates.

Parameters
Return type
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=10000, min_weight=None, uniform_weights=False)[source]

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, Ys, Yvars, parameter_values, outcome_names)[source]

Fit model to m outcomes.

Parameters
Return type

None

gen(n, parameter_values, objective_weights, outcome_constraints=None, fixed_features=None, pending_observations=None, model_gen_options=None)[source]

Generate new candidates.

Parameters
Return type
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)[source]

Predict

Parameters

X (List[List[Union[str, float, int, None]]]) – List of the j parameterizations at which to make predictions.

Return type

Tuple[ndarray, ndarray]

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

## NumPy Models¶

### ax.models.numpy.randomforest module¶

class ax.models.numpy.randomforest.RandomForest(max_features='sqrt', num_trees=500)[source]

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
cross_validate(Xs_train, Ys_train, Yvars_train, X_test)[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 (List[ndarray]) – A list of m (k_i x d) feature matrices X. Number of rows k_i can vary from i=1,…,m.

• Ys_train (List[ndarray]) – The corresponding list of m (k_i x 1) outcome arrays Y, for each outcome.

• Yvars_train (List[ndarray]) – The variances of each entry in Ys, same shape.

• X_test (ndarray) – (j x d) array of the j points at which to make predictions.

Return type

Tuple[ndarray, ndarray]

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, Ys, Yvars, bounds, task_features, feature_names, fidelity_features)[source]

Fit model to m outcomes.

Parameters
Return type

None

predict(X)[source]

Predict

Parameters

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

Return type

Tuple[ndarray, ndarray]

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=False, seed=None)[source]

Bases: ax.models.base.Model

This class specifies the basic skeleton for a random model.

As random generators do not make use of models, they do not implement the fit or predict methods.

These models do not need data, or optimization configs.

To satisfy search space parameter constraints, these models can use rejection sampling. To enable rejection sampling for a subclass, only only _gen_samples needs to be implemented, or alternatively, _gen_unconstrained/gen can be directly implemented.

deduplicate

If specified, 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, bounds, linear_constraints=None, fixed_features=None, model_gen_options=None, rounding_func=None)[source]

Generate new candidates.

Parameters
Return type

Tuple[ndarray, ndarray]

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.sobol module¶

class ax.models.random.sobol.SobolGenerator(seed=None, deduplicate=False, init_position=0, scramble=True)[source]

This class specifies the generation algorithm for a Sobol generator.

As Sobol does not make use of a model, it does not implement the fit or predict methods.

deduplicate

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

init_position

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

scramble

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

seed

An optional seed value for scrambling.

property engine

Return a singleton SobolEngine.

Return type

Optional[SobolEngine]

gen(n, bounds, linear_constraints=None, fixed_features=None, model_gen_options=None, rounding_func=None)[source]

Generate new candidates.

Parameters
Return type

Tuple[ndarray, ndarray]

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)[source]

Initialize singleton SobolEngine, only on gen.

Parameters

n_tunable_features (int) – The number of features which can be searched over.

Return type

SobolEngine

Returns

SobolEngine, which can generate Sobol points.

### ax.models.random.uniform module¶

class ax.models.random.uniform.UniformGenerator(deduplicate=False, seed=None)[source]

This class specifies a uniform random generation algorithm.

As a uniform generator does not make use of a model, it does not implement the fit or predict methods.

seed

An optional seed value for the underlying PRNG.

## Torch Models¶

### ax.models.torch.botorch module¶

class ax.models.torch.botorch.BotorchModel(model_constructor=<function get_and_fit_model>, model_predictor=<function predict_from_model>, acqf_constructor=<function get_NEI>, acqf_optimizer=<function scipy_optimizer>, refit_on_cv=False, refit_on_update=True, warm_start_refitting=True, **kwargs)[source]

Customizable botorch model.

By default, this uses a noisy Expected Improvement acquisition funciton 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 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

Call signatures:

model_constructor(
Xs,
Ys,
Yvars,
state_dict,
fidelity_features,
**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, state_dict is a pytorch module state dict, ‘fidelity_features’ is a list of ints that specify the positions of fidelity parameters in ‘Xs’, 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.

best_point(bounds, objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=None, model_gen_options=None)[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
Return type

Optional[Tensor]

Returns

d-tensor of the best point.

cross_validate(Xs_train, Ys_train, Yvars_train, X_test)[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 (List[Tensor]) – A list of m (k_i x d) feature tensors X. Number of rows k_i can vary from i=1,…,m.

• Ys_train (List[Tensor]) – The corresponding list of m (k_i x 1) outcome tensors Y, for each outcome.

• Yvars_train (List[Tensor]) – The variances of each entry in Ys, same shape.

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

Return type

Tuple[Tensor, Tensor]

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(Xs, Ys, Yvars, bounds, task_features, feature_names, fidelity_features)[source]

Fit model to m outcomes.

Parameters
Return type

None

gen(n, bounds, objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=None, pending_observations=None, model_gen_options=None, rounding_func=None)[source]

Generate new candidates.

An initialized acquisition function can be passed in as model_gen_options[“acquisition_function”].

Parameters
Returns

n x d-dim Tensor of generated points. Tensor: n-dim Tensor of weights for each point.

Return type

Tensor

predict(X)[source]

Predict

Parameters

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

Return type

Tuple[Tensor, Tensor]

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(Xs, Ys, Yvars)[source]

Update the model.

Updating the model requires both existing and additional data. The data passed into this method will become the new training data.

Parameters
• Xs (List[Tensor]) – Existing + additional data for the model, in the same format as for fit.

• Ys (List[Tensor]) – Existing + additional data for the model, in the same format as for fit.

• Yvars (List[Tensor]) – Existing + additional data for the model, in the same format as for fit.

Return type

None

ax.models.torch.botorch.get_rounding_func(rounding_func)[source]
Return type

Optional[Callable[[Tensor], Tensor]]

### ax.models.torch.utils module¶

ax.models.torch.utils.is_noiseless(model)[source]

Check if a given (single-task) botorch model is noiseless

Return type

bool