ax.generation_strategy

Generation Strategy

class ax.generation_strategy.generation_strategy.GenerationStrategy(steps: list[GenerationStep] | None = None, name: str | None = None, nodes: list[GenerationNode] | None = None)[source]

Bases: Base

GenerationStrategy describes which model should be used to generate new points for which trials, enabling and automating use of different models throughout the optimization process. For instance, it allows to use one model for the initialization trials, and another one for all subsequent trials. In the general case, this allows to automate use of an arbitrary number of models to generate an arbitrary numbers of trials described in the trials_per_model argument.

Parameters:

nodes – A list of GenerationNode. Each GenerationNode in the list represents a single node in a GenerationStrategy which, when composed of GenerationNodes, can be conceptualized as a graph instead of a linear list. TransitionCriterion defined in each GenerationNode represent the edges in the GenerationStrategy graph. GenerationNodes are more flexible than GenerationSteps and new GenerationStrategies should use nodes. Notably, either, but not both, of nodes and steps must be provided.
steps – A list of GenerationStep describing steps of this strategy.
name – An optional name for this generation strategy. If not specified, strategy’s name will be names of its nodes’ models joined with ‘+’.

DEFAULT_N: int = 1

clone_reset() → GenerationStrategy[source]: Copy this generation strategy without it’s state.

current_generator_run_limit() → tuple[int, bool][source]

First check if we can move the generation strategy to the next node, which is safe, as the next call to gen will just pick up from there. Then determine how many generator runs this generation strategy can generate right now, assuming each one of them becomes its own trial, and whether optimization is completed.

Returns: a two-item tuple of:

the number of generator runs that can currently be produced, with -1 meaning unlimited generator runs,
whether optimization is completed and the generation strategy cannot generate any more generator runs at all.

property current_node: GenerationNode: Current generation node.

property current_node_name: str: Current generation node name.

property current_step: GenerationStep: Current generation step.

property current_step_index: int: Returns the index of the current generation step. This attribute is replaced by node_name in newer GenerationStrategies but surfaced here for backward compatibility.

property experiment: Experiment: Experiment, currently set on this generation strategy.

gen(experiment: Experiment, data: Data | None = None, pending_observations: dict[str, list[ObservationFeatures]] | None = None, n: int = 1, fixed_features: ObservationFeatures | None = None) → GeneratorRun[source]

Produce the next points in the experiment. Additional kwargs passed to this method are propagated directly to the underlying model’s gen, along with the model_gen_kwargs set on the current generation node.

NOTE: Each generator run returned from this function must become a single trial on the experiment to comply with assumptions made in generation strategy. Do not split one generator run produced from generation strategy into multiple trials (never making a generator run into a trial is allowed).

Parameters:

experiment – Experiment, for which the generation strategy is producing a new generator run in the course of gen, and to which that generator run will be added as trial(s). Information stored on the experiment (e.g., trial statuses) is used to determine which model will be used to produce the generator run returned from this method.
data – Optional data to be passed to the underlying model’s gen, which is called within this method and actually produces the resulting generator run. By default, data is all data on the experiment.
n – Integer representing how many arms should be in the generator run produced by this method. NOTE: Some underlying models may ignore the n and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that can differ from n.
pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.

gen_for_multiple_trials_with_multiple_models(experiment: Experiment, data: Data | None = None, pending_observations: dict[str, list[ObservationFeatures]] | None = None, n: int | None = None, fixed_features: ObservationFeatures | None = None, num_trials: int = 1, arms_per_node: dict[str, int] | None = None) → list[list[GeneratorRun]][source]

Produce GeneratorRuns for multiple trials at once with the possibility of using multiple models per trial, getting multiple GeneratorRuns per trial.

Parameters:

experiment – Experiment, for which the generation strategy is producing a new generator run in the course of gen, and to which that generator run will be added as trial(s). Information stored on the experiment (e.g., trial statuses) is used to determine which model will be used to produce the generator run returned from this method.
data – Optional data to be passed to the underlying model’s gen, which is called within this method and actually produces the resulting generator run. By default, data is all data on the experiment.
pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.
n – Integer representing how many total arms should be in the generator runs produced by this method. NOTE: Some underlying models may ignore the n and produce a model-determined number of arms. In that case this method will also output generator runs with number of arms that can differ from n.
fixed_features – An optional set of ObservationFeatures that will be passed down to the underlying models. Note: if provided this will override any algorithmically determined fixed features so it is important to specify all necessary fixed features.
num_trials – Number of trials to generate generator runs for in this call. If not provided, defaults to 1.
arms_per_node – An optional map from node name to the number of arms to generate from that node. If not provided, will default to the number of arms specified in the node’s InputConstructors or n if no InputConstructors are defined on the node. We expect either n or arms_per_node to be provided, but not both, and this is an advanced argument that should only be used by advanced users.

Returns:

A list of lists of lists generator runs. Each outer list represents a trial being suggested and each inner list represents a generator run for that trial.

property is_node_based: bool: Whether this strategy consists of GenerationNodes only. This is useful for determining initialization properties and other logic.

property last_generator_run: GeneratorRun | None: Latest generator run produced by this generation strategy. Returns None if no generator runs have been produced yet.

property model: Adapter | None: Current model in this strategy. Returns None if no model has been set yet (i.e., if no generator runs have been produced from this GS).

property model_transitions: list[int]: [DEPRECATED]List of trial indices where a transition happened from one model to another.

property name: str: Name of this generation strategy.

property nodes_dict: dict[str, GenerationNode]: Returns a dictionary mapping node names to nodes.

property optimization_complete: bool: Checks whether all nodes are completed in the generation strategy.

property trials_as_df: DataFrame | None

Puts information on individual trials into a data frame for easy viewing.

THIS METHOD IS DEPRECATED AND WILL BE REMOVED IN A FUTURE RELEASE. Please use Experiment.to_df() instead.

property uses_non_registered_models: bool: Whether this generation strategy involves models that are not registered and therefore cannot be stored.

ax.generation_strategy.generation_strategy.step_based_gs_only(f: Callable[[...], T]) → Callable[[...], T][source]: For use as a decorator on functions only implemented for GenerationStep-based GenerationStrategies. Mainly useful for older GenerationStrategies.

Generation Node

class ax.generation_strategy.generation_node.GenerationNode(node_name: str, model_specs: list[GeneratorSpec], best_model_selector: BestModelSelector | None = None, should_deduplicate: bool = False, transition_criteria: Sequence[TransitionCriterion] | None = None, input_constructors: None | dict[InputConstructorPurpose, NodeInputConstructors] = None, previous_node_name: str | None = None, trial_type: str | None = None, should_skip: bool = False, fallback_specs: dict[type[Exception], GeneratorSpec] | None = None)[source]

Bases: SerializationMixin, SortableBase

Base class for GenerationNode, capable of fitting one or more model specs under the hood and generating candidates from them.

Parameters:

node_name – A unique name for the GenerationNode. Used for storage purposes.
model_specs – A list of GeneratorSpecs to be selected from for generation in this GenerationNode.
best_model_selector – A BestModelSelector used to select the GeneratorSpec to generate from in GenerationNode with multiple ``GeneratorSpec``s.
should_deduplicate – Whether to deduplicate the parameters of proposed arms against those of previous arms via rejection sampling. If this is True, the GenerationStrategy will discard generator runs produced from the GenerationNode that has should_deduplicate=True if they contain arms already present on the experiment and replace them with new generator runs. If no generator run with entirely unique arms could be produced in 5 attempts, a GenerationStrategyRepeatedPoints error will be raised, as we assume that the optimization converged when the model can no longer suggest unique arms.
transition_criteria – List of TransitionCriterion, each of which describes a condition that must be met before completing a GenerationNode. All is_met must evaluateTrue for the GenerationStrategy to move on to the next GenerationNode.
input_constructors – A dictionary mapping input constructor purpose enum to the input constructor enum. Each input constructor maps to a method which encodes the logic for determining dynamic inputs to the GenerationNode
trial_type – Specifies the type of trial to generate, is limited to either Keys.SHORT_RUN or Keys.LONG_RUN for now. If not specified, will default to None and not be used during generation.
previous_node_name – The previous GenerationNode name in the GenerationStrategy, if any. Initialized to None for all nodes, and is set during transition from one GenerationNode to the next. Can be overwritten if multiple transitions occur between nodes, and will always store the most recent previous GenerationNode name.
should_skip – Whether to skip this node during generation time. Defaults to False, and can only currently be set to True via NodeInputConstructors
fallback_specs – Optional dict mapping expected exception types to ModelSpec fallbacks used when gen fails.

Note for developers: by “model” here we really mean an Ax Adapter object, which contains an Ax Model under the hood. We call it “model” here to simplify and focus on explaining the logic of GenerationStep and GenerationStrategy.

apply_input_constructors(experiment: Experiment, gen_kwargs: dict[str, Any]) → dict[str, int | ObservationFeatures | None][source]

property experiment: Experiment: Returns the experiment associated with this GenerationStrategy

fallback_specs: dict[type[Exception], GeneratorSpec]

gen(*, experiment: Experiment, pending_observations: dict[str, list[ObservationFeatures]] | None, skip_fit: bool = False, data: Data | None = None, **gs_gen_kwargs: Any) → GeneratorRun | None[source]

This method generates candidates using self._gen and handles deduplication of generated candidates if self.should_deduplicate=True.

NOTE: This method returning None indicates that this node should be skipped: by its generation strategy.
NOTE: n argument will be passed through this node’s input constructors,: which might modify the number of arms this node will produce. Also, some underlying generators may ignore the n argument and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that may differ from n.

Parameters:

experiment – The experiment to generate candidates for.
data – Optional override for the experiment data used to generate candidates; if not specified, will use experiment.lookup_data() (extracted in Adapter).
pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.
model_gen_kwargs – Keyword arguments, passed through to ModelSpec.gen; these override any pre-specified in ModelSpec.model_gen_kwargs. Often will contain n.

Returns:

A GeneratorRun containing the newly generated candidates or None if this node is not in a correct state to generate candidates and should be skipped (e.g. if its input constructor for the n argument specifies that it should generate 0 candidate arms given the current experiment state and user inputs).

property generation_strategy: GenerationStrategy: Returns a backpointer to the GenerationStrategy, useful for obtaining the experiment associated with this GenerationStrategy

property input_constructors: dict[InputConstructorPurpose, NodeInputConstructors]: Returns the input constructors that will be used to determine any dynamic inputs to this GenerationNode.

property is_completed: bool: Returns True if this GenerationNode is complete and should transition to the next node.

property model_spec_to_gen_from: GeneratorSpec: Returns the cached _model_spec_to_gen_from or gets it from _pick_fitted_model_to_gen_from and then caches and returns it

model_specs: list[GeneratorSpec]

property model_to_gen_from_name: str | None: Returns the name of the model that will be used for gen, if there is one. Otherwise, returns None.

new_trial_limit(raise_generation_errors: bool = False) → int[source]

How many trials can this generation strategy can currently produce GeneratorRun-s for (with potentially multiple generator runs produced for each intended trial).

NOTE: Only considers transition criteria that inherit from TrialBasedCriterion.

Returns:: The number of generator runs that can currently be produced, with -1 meaning unlimited generator runs.

property node_name: str: Returns the unique name of this GenerationNode

property node_that_generated_last_gr: str | None

Returns the name of the node that generated the last generator run.

Returns:: The name of the node that generated the last generator run.
Return type:: str

property previous_node: GenerationNode | None: Returns the previous GenerationNode, if any.

should_deduplicate: bool

should_transition_to_next_node(raise_data_required_error: bool = True) → tuple[bool, str][source]

Checks whether we should transition to the next node based on this node’s TransitionCriterion.

Important: This method relies on the transition_criterion of this node to be listed in order of importance. Ex: a fallback transition should come after the primary transition in the transition criterion list.

Parameters:

raise_data_required_error – Whether to raise DataRequiredError in the case detailed above. Not raising the error is useful if just looking to check how many generator runs (to be made into trials) can be produced, but not actually producing them yet.

Returns:

Whether we should transition to the next node: and the name of the node to gen from (either the current or next node)

Return type:

Tuple[bool, str]

property transition_criteria: Sequence[TransitionCriterion]: Returns the sequence of TransitionCriteria that will be used to determine if this GenerationNode is complete and should transition to the next node.

property transition_edges: dict[str, list[TransitionCriterion]]

Returns a dictionary mapping the next GenerationNode to the TransitionCriteria that define the transition that that node.

Ex: if the transition from the current node to node x is defined by IsSingleObjective and MinTrials criterion then the return would be {‘x’: [IsSingleObjective, MinTrials]}.

Returns:: A dictionary mapping the next GenerationNode to the TransitionCriterion that are associated with it.
Return type:: Dict[str, List[TransitionCriterion]]

property trials_from_node: set[int]

Returns a set mapping a GenerationNode to the trials it generated.

Returns:: A set containing all the trials indices generated by this node.
Return type:: Set[int]

class ax.generation_strategy.generation_node.GenerationStep(model: ModelRegistryBase | Callable[[...], Adapter], num_trials: int, model_kwargs: dict[str, Any] | None = None, model_gen_kwargs: dict[str, Any] | None = None, completion_criteria: Sequence[TransitionCriterion] | None = None, min_trials_observed: int = 0, max_parallelism: int | None = None, enforce_num_trials: bool = True, should_deduplicate: bool = False, model_name: str | None = None, use_update: bool = False, index: int = -1)[source]

Bases: GenerationNode, SortableBase

One step in the generation strategy, corresponds to a single model. Describes the model, how many trials will be generated with this model, what minimum number of observations is required to proceed to the next model, etc.

NOTE: Model can be specified either from the model registry (ax.modelbridge.registry.Generators or using a callable model constructor. Only models from the registry can be saved, and thus optimization can only be resumed if interrupted when using models from the registry.

Parameters:

model – A member of Generators enum or a callable returning an instance of Adapter with an instantiated underlying Model. Refer to ax/modelbridge/factory.py for examples of such callables.
num_trials – How many trials to generate with the model from this step. If set to -1, trials will continue to be generated from this model as long as generation_strategy.gen is called (available only for the last of the generation steps).
min_trials_observed – How many trials must be completed before the generation strategy can proceed to the next step. Defaults to 0. If num_trials of a given step have been generated but min_trials_ observed have not been completed, a call to generation_strategy.gen will fail with a DataRequiredError.
max_parallelism – How many trials generated in the course of this step are allowed to be run (i.e. have trial.status of RUNNING) simultaneously. If max_parallelism trials from this step are already running, a call to generation_strategy.gen will fail with a MaxParallelismReached Exception, indicating that more trials need to be completed before generating and running next trials.
use_update – DEPRECATED.
enforce_num_trials – Whether to enforce that only num_trials are generated from the given step. If False and num_trials have been generated, but min_trials_observed have not been completed, generation_strategy.gen will continue generating trials from the current step, exceeding num_ trials for it. Allows to avoid DataRequiredError, but delays proceeding to next generation step.
model_kwargs – Dictionary of kwargs to pass into the model constructor on instantiation. E.g. if model is Generators.SOBOL, kwargs will be applied as Generators.SOBOL(**model_kwargs); if model is get_sobol, get_sobol(**model_kwargs). NOTE: if generation strategy is interrupted and resumed from a stored snapshot and its last used model has state saved on its generator runs, model_kwargs is updated with the state dict of the model, retrieved from the last generator run of this generation strategy.
model_gen_kwargs – Each call to generation_strategy.gen performs a call to the step’s model’s gen under the hood; model_gen_kwargs will be passed to the model’s gen like so: model.gen(**model_gen_kwargs).
completion_criteria – List of TransitionCriterion. All is_met must evaluate True for the GenerationStrategy to move on to the next Step
index – Index of this generation step, for use internally in Generation Strategy. Do not assign as it will be reassigned when instantiating GenerationStrategy with a list of its steps.
should_deduplicate – Whether to deduplicate the parameters of proposed arms against those of previous arms via rejection sampling. If this is True, the generation strategy will discard generator runs produced from the generation step that has should_deduplicate=True if they contain arms already present on the experiment and replace them with new generator runs. If no generator run with entirely unique arms could be produced in 5 attempts, a GenerationStrategyRepeatedPoints error will be raised, as we assume that the optimization converged when the model can no longer suggest unique arms.
model_name – Optional name of the model. If not specified, defaults to the model key of the model spec.

Note for developers: by “model” here we really mean an Ax Adapter object, which contains an Ax Model under the hood. We call it “model” here to simplify and focus on explaining the logic of GenerationStep and GenerationStrategy.

gen(*, experiment: Experiment, n: int | None = None, pending_observations: dict[str, list[ObservationFeatures]] | None, skip_fit: bool = False, data: Data | None = None, **gs_gen_kwargs: Any) → GeneratorRun | None[source]

This method generates candidates using self._gen and handles deduplication of generated candidates if self.should_deduplicate=True.

NOTE: This method returning None indicates that this node should be skipped: by its generation strategy.
NOTE: n argument will be passed through this node’s input constructors,: which might modify the number of arms this node will produce. Also, some underlying generators may ignore the n argument and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that may differ from n.

Parameters:

experiment – The experiment to generate candidates for.
data – Optional override for the experiment data used to generate candidates; if not specified, will use experiment.lookup_data() (extracted in Adapter).
pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.
model_gen_kwargs – Keyword arguments, passed through to ModelSpec.gen; these override any pre-specified in ModelSpec.model_gen_kwargs. Often will contain n.

Returns:

A GeneratorRun containing the newly generated candidates or None if this node is not in a correct state to generate candidates and should be skipped (e.g. if its input constructor for the n argument specifies that it should generate 0 candidate arms given the current experiment state and user inputs).

property model_gen_kwargs: dict[str, Any]: Returns the model gen kwargs of the underlying GeneratorSpec.

property model_kwargs: dict[str, Any]: Returns the model kwargs of the underlying GeneratorSpec.

property model_spec: GeneratorSpec: Returns the first model_spec from the model_specs attribute.

External Generation Node

class ax.generation_strategy.external_generation_node.ExternalGenerationNode(node_name: str, should_deduplicate: bool = True, transition_criteria: Sequence[TransitionCriterion] | None = None)[source]

Bases: GenerationNode, ABC

A generation node intended to be used with non-Ax methods for candidate generation.

To leverage external methods for candidate generation, the user must create a subclass that implements update_generator_state and get_next_candidate methods. This can then be provided as a node into a GenerationStrategy, either as standalone or as part of a larger generation strategy with other generation nodes, e.g., with a Sobol node for initialization.

Example: >>> class MyExternalGenerationNode(ExternalGenerationNode): >>> … >>> generation_strategy = GenerationStrategy( >>> nodes = [MyExternalGenerationNode(…)] >>> ) >>> ax_client = AxClient(generation_strategy=generation_strategy) >>> ax_client.create_experiment(…) >>> ax_client.get_next_trial() # Generates trials using the new generation node.

Get the parameters for the next candidate configuration to evaluate.

Parameters:: pending_parameters – A list of parameters of the candidates pending evaluation. This is often used to avoid generating duplicate candidates.
Returns:: A dictionary mapping parameter names to parameter values for the next candidate suggested by the method.

property model_spec_to_gen_from: None: Returns the cached _model_spec_to_gen_from or gets it from _pick_fitted_model_to_gen_from and then caches and returns it

abstractmethod update_generator_state(experiment: Experiment, data: Data) → None[source]

A method used to update the state of the generator. This includes any models, predictors or any other custom state used by the generation node. This method will be called with the up-to-date experiment and data before get_next_candidate is called to generate the next trial(s). Note that get_next_candidate may be called multiple times (to generate multiple candidates) after a call to update_generator_state.

Parameters:

experiment – The Experiment object representing the current state of the experiment. The key properties includes trials, search_space, and optimization_config. The data is provided as a separate arg.
data – The data / metrics collected on the experiment so far.

Transition Criterion

class ax.generation_strategy.transition_criterion.AutoTransitionAfterGen(transition_to: str, block_transition_if_unmet: bool | None = True, continue_trial_generation: bool | None = True)[source]

Bases: TransitionCriterion

A class to designate automatic transition from one GenerationNode to another.

Parameters:

transition_to – The name of the GenerationNode the GenerationStrategy should transition to next.
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: This criterion defaults to setting this to True to ensure we validate a GeneratorRun is generated by the current GenerationNode.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial.

block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Error to be raised if the block_gen_if_met flag is set to True.

is_met(experiment: Experiment, curr_node: GenerationNode) → bool[source]: Return True as soon as any GeneratorRun is generated by this GenerationNode.

class ax.generation_strategy.transition_criterion.AuxiliaryExperimentCheck(transition_to: str, auxiliary_experiment_purposes_to_include: list[AuxiliaryExperimentPurpose] | None = None, auxiliary_experiment_purposes_to_exclude: list[AuxiliaryExperimentPurpose] | None = None, block_transition_if_unmet: bool | None = True, block_gen_if_met: bool | None = False, continue_trial_generation: bool | None = False)[source]

Bases: TransitionCriterion

A class to transition from one GenerationNode to another by checking if certain types of Auxiliary Experiment purposes exists.

A common use case is to use auxiliary_experiment_purposes_to_include to transition to a node and auxiliary_experiment_purposes_to_exclude to transition away from it.

Example usage: In Bayesian optimization with preference exploration (BOPE), we check if the preference exploration (PE) auxiliary experiment exists to indicate transition to the node that will generate candidates based on the learned objective.Since preference exploration is usually conducted after the exploratory batch is completed, we do not know at experiment creation time if the PE node should be used during the GenerationStrategy.

Parameters:

transition_to – The name of the GenerationNode the GenerationStrategy should transition to when this criterion is met, if it exists.
auxiliary_experiment_purposes_to_include – Optional list of auxiliary experiment purposes we expect to have. This can be helpful when need to transition to a node based on AuxiliaryExperimentPurpose. Criterion is met when all inclusion and exclusion checks pass.
auxiliary_experiment_purposes_to_exclude – Optional list of auxiliary experiment purpose we expect to not have. This can be helpful when need to transition out of a node based on AuxiliaryExperimentPurpose. Criterion is met when all inclusion and exclusion checks pass.
block_gen_if_met – A flag to prevent continued generation from the associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinTrials has not been met yet, but MinTrials has been reached. If this flag is set to true on MinTrials then we will raise an error, otherwise we will continue to generate trials until MinTrials is met (thus overriding MinTrials).
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: MaxGenerationParallelism defaults to setting this to False since we can complete and move on from this node without ever reaching its threshold.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial.

block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Raises the appropriate error (should only be called when the GenerationNode is blocked from continued generation). For this class, the exception is DataRequiredError.

check_aux_exp_purposes(aux_exp_purposes: Collection[AuxiliaryExperimentPurpose], include: bool, expected_aux_exp_purposes: list[AuxiliaryExperimentPurpose] | None = None) → bool[source]: Helper method to check if all elements in expected_aux_exp_purposes are in (or not in) aux_exp_purposes

is_met(experiment: Experiment, curr_node: GenerationNode) → bool[source]: Check if the experiment has auxiliary experiments for certain purpose.

class ax.generation_strategy.transition_criterion.IsSingleObjective(transition_to: str, block_transition_if_unmet: bool | None = True, continue_trial_generation: bool | None = False)[source]

Bases: TransitionCriterion

A class to initiate transition based on whether the experiment is optimizing for a single objective or multiple objectives.

Parameters:

transition_to – The name of the GenerationNode the GenerationStrategy should transition to next.
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: This criterion defaults to setting this to True to ensure we validate a GeneratorRun is generated by the current GenerationNode.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial.

block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Error to be raised if the block_gen_if_met flag is set to True.

is_met(experiment: Experiment, curr_node: GenerationNode) → bool[source]: Return True if the optimization config is not of type MultiObjectiveOptimizationConfig.

class ax.generation_strategy.transition_criterion.MaxGenerationParallelism(threshold: int, only_in_statuses: list[TrialStatus] | None = None, not_in_statuses: list[TrialStatus] | None = None, transition_to: str | None = None, block_transition_if_unmet: bool | None = False, block_gen_if_met: bool | None = True, use_all_trials_in_exp: bool | None = False, continue_trial_generation: bool | None = True)[source]

Bases: TrialBasedCriterion

Specific TransitionCriterion implementation which defines the maximum number of trials that can simultaneously be in the designated trial statuses. The default behavior is to block generation from the associated GenerationNode if the threshold is met. This is configured via the block_gen_if_met flag being set to True. This criterion defaults to not blocking transition to another node via the block_transition_if_unmet flag being set to False.

Parameters:

threshold – The threshold as an integer for this criterion. Ex: If we want to generate at most 3 trials, then the threshold is 3.
only_in_statuses – A list of trial statuses to filter on when checking the criterion threshold.
not_in_statuses – A list of trial statuses to exclude when checking the criterion threshold.
transition_to – The name of the GenerationNode the GenerationStrategy should transition to when this criterion is met, if it exists.
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: MaxGenerationParallelism defaults to setting this to False since we can complete and move on from this node without ever reaching its threshold.
block_gen_if_met – A flag to prevent continued generation from the associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinTrials has not been met yet, but MinTrials has been reached. If this flag is set to true on MinTrials then we will raise an error, otherwise we will continue to generate trials until MinTrials is met (thus overriding MinTrials).
use_all_trials_in_exp – A flag to use all trials in the experiment, instead of only those generated by the current GenerationNode.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial. Defaults to False for MaxGenerationParallelism since this criterion isn’t currently used for node -> node or trial -> trial transition.
count_only_trials_with_data – If set to True, only trials with data will be counted towards the threshold. Defaults to False.

block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Raises the appropriate error (should only be called when the GenerationNode is blocked from continued generation). For this class, the exception is MaxParallelismReachedException.

class ax.generation_strategy.transition_criterion.MaxTrials(threshold: int, only_in_statuses: list[TrialStatus] | None = None, not_in_statuses: list[TrialStatus] | None = None, transition_to: str | None = None, block_transition_if_unmet: bool | None = True, block_gen_if_met: bool | None = False, use_all_trials_in_exp: bool | None = False, continue_trial_generation: bool | None = False, count_only_trials_with_data: bool = False)[source]

Bases: TrialBasedCriterion

Simple class to enforce a maximum threshold for the number of trials with the designated statuses being generated by a specific GenerationNode. The default behavior is to block transition to the next node if the threshold is unmet, but not affect continued generation.

Parameters:

threshold – The threshold as an integer for this criterion. Ex: If we want to generate at most 3 trials, then the threshold is 3.
only_in_statuses – A list of trial statuses to filter on when checking the criterion threshold.
not_in_statuses – A list of trial statuses to exclude when checking the criterion threshold.
transition_to – The name of the GenerationNode the GenerationStrategy should transition to when this criterion is met, if it exists.
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: MaxGenerationParallelism defaults to setting this to False since we can complete and move on from this node without ever reaching its threshold.
block_gen_if_met – A flag to prevent continued generation from the associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinimumTrialsInStatus has not been met yet, but MaxTrials has been reached. If this flag is set to true on MaxTrials then we will raise an error, otherwise we will continue to generate trials until MinimumTrialsInStatus is met (thus overriding MaxTrials).
use_all_trials_in_exp – A flag to use all trials in the experiment, instead of only those generated by the current GenerationNode.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial.
count_only_trials_with_data – If set to True, only trials with data will be counted towards the threshold. Defaults to False.

block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Raises the appropriate error (should only be called when the GenerationNode is blocked from continued generation). For this class, the exception is DataRequiredError.

class ax.generation_strategy.transition_criterion.MinTrials(threshold: int, only_in_statuses: list[TrialStatus] | None = None, not_in_statuses: list[TrialStatus] | None = None, transition_to: str | None = None, block_transition_if_unmet: bool | None = True, block_gen_if_met: bool | None = False, use_all_trials_in_exp: bool | None = False, continue_trial_generation: bool | None = False, count_only_trials_with_data: bool = False)[source]

Bases: TrialBasedCriterion

Simple class to enforce a minimum threshold for the number of trials with the designated statuses being generated by a specific GenerationNode. The default behavior is to block transition to the next node if the threshold is unmet, but not affect continued generation.

Parameters:

threshold – The threshold as an integer for this criterion. Ex: If we want to generate at most 3 trials, then the threshold is 3.
only_in_statuses – A list of trial statuses to filter on when checking the criterion threshold.
not_in_statuses – A list of trial statuses to exclude when checking the criterion threshold.
transition_to – The name of the GenerationNode the GenerationStrategy should transition to when this criterion is met, if it exists.
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: MaxGenerationParallelism defaults to setting this to False since we can complete and move on from this node without ever reaching its threshold.
block_gen_if_met – A flag to prevent continued generation from the associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinTrials has not been met yet, but MinTrials has been reached. If this flag is set to true on MinTrials then we will raise an error, otherwise we will continue to generate trials until MinTrials is met (thus overriding MinTrials).
use_all_trials_in_exp – A flag to use all trials in the experiment, instead of only those generated by the current GenerationNode.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial.
count_only_trials_with_data – If set to True, only trials with data will be counted towards the threshold. Defaults to False.

block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Raises the appropriate error (should only be called when the GenerationNode is blocked from continued generation). For this class, the exception is DataRequiredError.

class ax.generation_strategy.transition_criterion.MinimumPreferenceOccurances(metric_name: str, threshold: int, transition_to: str | None = None, block_gen_if_met: bool | None = False, block_transition_if_unmet: bool | None = True)[source]

Bases: TransitionCriterion

In a preference Experiment (i.e. Metric values may either be zero for No and nonzero for Yes) do not transition until a minimum number of both Yes and No responses have been received.

Parameters:

metric_name – name of the metric to check for preference occurrences.
threshold – The threshold as an integer for this criterion. Ex: If we want to generate at most 3 trials, then the threshold is 3.
transition_to – The name of the GenerationNode the GenerationStrategy should transition to when this criterion is met, if it exists.
block_gen_if_met – A flag to prevent continued generation from the associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinTrials has not been met yet, but MinTrials has been reached. If this flag is set to true on MinTrials then we will raise an error, otherwise we will continue to generate trials until MinTrials is met (thus overriding MinTrials).
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: MaxGenerationParallelism defaults to setting this to False since we can complete and move on from this node without ever reaching its threshold.

block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Error to be raised if the block_gen_if_met flag is set to True.

is_met(experiment: Experiment, curr_node: GenerationNode) → bool[source]: If the criterion of this TransitionCriterion is met, returns True.

class ax.generation_strategy.transition_criterion.MinimumTrialsInStatus(status: TrialStatus, threshold: int, transition_to: str | None = None)[source]

Bases: TransitionCriterion

Deprecated and replaced with more flexible MinTrials criterion.

block_continued_generation_error(node_name: str | None, model_name: str | None, experiment: Experiment | None, trials_from_node: set[int]) → None[source]: Error to be raised if the block_gen_if_met flag is set to True.

is_met(experiment: Experiment, curr_node: GenerationNode) → bool[source]: If the criterion of this TransitionCriterion is met, returns True.

class ax.generation_strategy.transition_criterion.TransitionCriterion(transition_to: str | None = None, block_transition_if_unmet: bool | None = True, block_gen_if_met: bool | None = False, continue_trial_generation: bool | None = False)[source]

Bases: SortableBase, SerializationMixin

Simple class to describe a condition which must be met for this GenerationNode to take an action such as generation, transition, etc.

Parameters:

transition_to – The name of the GenerationNode the GenerationStrategy should transition to when this criterion is met, if it exists.
block_gen_if_met – A flag to prevent continued generation from the associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinTrials has not been met yet, but MinTrials has been reached. If this flag is set to true on MinTrials then we will raise an error, otherwise we will continue to generate trials until MinTrials is met (thus overriding MinTrials).
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: MaxGenerationParallelism defaults to setting this to False since we can complete and move on from this node without ever reaching its threshold.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial.

abstractmethod block_continued_generation_error(node_name: str, experiment: Experiment, trials_from_node: set[int]) → None[source]: Error to be raised if the block_gen_if_met flag is set to True.

property criterion_class: str: Name of the class of this TransitionCriterion.

abstractmethod is_met(experiment: Experiment, curr_node: GenerationNode) → bool[source]: If the criterion of this TransitionCriterion is met, returns True.

property transition_to: str | None: The name of the next GenerationNode after this TransitionCriterion is completed, if it exists.

class ax.generation_strategy.transition_criterion.TrialBasedCriterion(threshold: int, block_transition_if_unmet: bool | None = True, block_gen_if_met: bool | None = False, only_in_statuses: list[TrialStatus] | None = None, not_in_statuses: list[TrialStatus] | None = None, transition_to: str | None = None, use_all_trials_in_exp: bool | None = False, continue_trial_generation: bool | None = False, count_only_trials_with_data: bool = False)[source]

Bases: TransitionCriterion

Common class for transition criterion that are based on trial information.

Parameters:

threshold – The threshold as an integer for this criterion. Ex: If we want to generate at most 3 trials, then the threshold is 3.
block_transition_if_unmet – A flag to prevent the node from completing and being able to transition to another node. Ex: MaxGenerationParallelism defaults to setting this to False since we can complete and move on from this node without ever reaching its threshold.
block_gen_if_met – A flag to prevent continued generation from the associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinTrials has not been met yet, but MinTrials has been reached. If this flag is set to true on MinTrials then we will raise an error, otherwise we will continue to generate trials until MinTrials is met (thus overriding MinTrials).
only_in_statuses – A list of trial statuses to filter on when checking the criterion threshold.
not_in_statuses – A list of trial statuses to exclude when checking the criterion threshold.
transition_to – The name of the GenerationNode the GenerationStrategy should transition to when this criterion is met, if it exists.
use_all_trials_in_exp – A flag to use all trials in the experiment, instead of only those generated by the current GenerationNode.
complete_trial_generation – A flag to indicate that all generation for a given trial is completed. This is necessary because in BatchTrial there are multiple arms per trial, and we enable generation of arms within a batch from different GenerationNodes. This flag should be set to True for the last node in a set of GenerationNodes expected to create a given BatchTrial.
count_only_trials_with_data – If set to True, only trials with data will be counted towards the threshold. Defaults to False.

all_trials_to_check(experiment: Experiment) → set[int][source]

All the trials to check from the entire experiment that meet all the provided status filters.

Parameters:: experiment – The experiment associated with this GenerationStrategy.

experiment_trials_by_status(experiment: Experiment, statuses: list[TrialStatus]) → set[int][source]

Get the trial indices from the entire experiment with the desired statuses.

Parameters:

experiment – The experiment associated with this GenerationStrategy.
statuses – The trial statuses to filter on.

Returns:

The trial indices in the experiment with the desired statuses.

is_met(experiment: Experiment, curr_node: GenerationNode) → bool[source]: Returns if this criterion has been met given its constraints. :param experiment: The experiment associated with this GenerationStrategy. :param trials_from_node: The set of trials generated by this GenerationNode. :param block_continued_generation: A flag to prevent continued generation from the

associated GenerationNode if this criterion is met but other criterion remain unmet. Ex: MinTrials has not been met yet, but MinTrials has been reached. If this flag is set to true on MinTrials then we will raise an error, otherwise we will continue to generate trials until MinTrials is met (thus overriding MinTrials).

num_contributing_to_threshold(experiment: Experiment, trials_from_node: set[int]) → int[source]

Returns the number of trials contributing to the threshold.

Parameters:

experiment – The experiment associated with this GenerationStrategy.
trials_from_node – The set of trials generated by this GenerationNode.

num_till_threshold(experiment: Experiment, trials_from_node: set[int]) → int[source]

Returns the number of trials needed to meet the threshold.

Parameters:

experiment – The experiment associated with this GenerationStrategy.
trials_from_node – The set of trials generated by this GenerationNode.

Generation Node Input Constructors

class ax.generation_strategy.generation_node_input_constructors.InputConstructorPurpose(*values)[source]

Bases: Enum

A simple enum to indicate the purpose of the input constructor. Values in this enum will match argument names for GenerationNode.gen.

Explanation of the different purposes:

N: Defines the logic to determine the number of arms to generate from the: next GenerationNode given the total number of arms expected in this trial.

FIXED_FEATURES = 'fixed_features'

N = 'n'

class ax.generation_strategy.generation_node_input_constructors.NodeInputConstructors(*values)[source]

Bases: FuncEnum

An enum which maps to a the name of a callable method for constructing GenerationNode inputs. Recommendation: ends of the names of members of this enum should match the corresponding InputConstructorPurpose name.

NOTE: All functions defined by this enum should share identical arguments in their signatures, and the return type should be the same across all functions that are used for the same InputConstructorPurpose.

ALL_N = 'consume_all_n'

REMAINING_N = 'remaining_n'

REPEAT_N = 'repeat_arm_n'

TARGET_TRIAL_FIXED_FEATURES = 'set_target_trial'

ax.generation_strategy.generation_node_input_constructors.consume_all_n(previous_node: GenerationNode | None, next_node: GenerationNode, gs_gen_call_kwargs: dict[str, Any], experiment: Experiment) → int[source]

Generate total requested number of arms from the next node.

Example: Initial exploration with Sobol will generate all arms from a single sobol node.

Note: If no n is provided to the GenerationStrategy gen call, we will use the default number of arms for the next node, defined as a constant DEFAULT_N in the GenerationStrategy file.

Parameters:

previous_node – The previous node in the GenerationStrategy. This is the node that is being transition away from, and is provided for easy access to properties of this node.
next_node – The next node in the GenerationStrategy. This is the node that will leverage the inputs defined by this input constructor.
gs_gen_call_kwargs – The kwargs passed to the GenerationStrategy’s gen call.
experiment – The experiment associated with this GenerationStrategy.

Returns:

The total number of requested arms from the next node.

ax.generation_strategy.generation_node_input_constructors.remaining_n(previous_node: GenerationNode | None, next_node: GenerationNode, gs_gen_call_kwargs: dict[str, Any], experiment: Experiment) → int[source]

Generate the remaining number of arms requested for this trial in gs.gen().

Note: If no n is provided to the GenerationStrategy gen call, we will use the default number of arms for the next node, defined as a constant DEFAULT_N in the GenerationStrategy file.

Parameters:

previous_node – The previous node in the GenerationStrategy. This is the node that is being transition away from, and is provided for easy access to properties of this node.
next_node – The next node in the GenerationStrategy. This is the node that will leverage the inputs defined by this input constructor.
gs_gen_call_kwargs – The kwargs passed to the GenerationStrategy’s gen call.
experiment – The experiment associated with this GenerationStrategy.

Returns:

The number of requested arms from the next node

ax.generation_strategy.generation_node_input_constructors.repeat_arm_n(previous_node: GenerationNode | None, next_node: GenerationNode, gs_gen_call_kwargs: dict[str, Any], experiment: Experiment) → int[source]

Generate a small percentage of arms requested to be used for repeat arms in the next trial.

Note: If no n is provided to the GenerationStrategy gen call, we will use the default number of arms for the next node, defined as a constant DEFAULT_N in the GenerationStrategy file.

Parameters:

previous_node – The previous node in the GenerationStrategy. This is the node that is being transition away from, and is provided for easy access to properties of this node.
next_node – The next node in the GenerationStrategy. This is the node that will leverage the inputs defined by this input constructor.
gs_gen_call_kwargs – The kwargs passed to the GenerationStrategy’s gen call.
experiment – The experiment associated with this GenerationStrategy.

Returns:

The number of requested arms from the next node

ax.generation_strategy.generation_node_input_constructors.set_target_trial(previous_node: GenerationNode | None, next_node: GenerationNode, gs_gen_call_kwargs: dict[str, Any], experiment: Experiment) → ObservationFeatures | None[source]

Determine the target trial for the next node based on the current state of the Experiment.

Args:

previous_node: The previous node in the GenerationStrategy. This is the node
that is being transition away from, and is provided for easy access to properties of this node.

next_node: The next node in the GenerationStrategy. This is the node that
will leverage the inputs defined by this input constructor.

gs_gen_call_kwargs: The kwargs passed to the GenerationStrategy’s
gen call.

experiment: The experiment associated with this GenerationStrategy.

Returns:: An ObservationFeatures object that defines the target trial for the next node.

GeneratorSpec

class ax.generation_strategy.model_spec.FactoryFunctionGeneratorSpec(model_enum: 'ModelRegistryBase | None' = None, model_kwargs: 'dict[str, Any]' = <factory>, model_gen_kwargs: 'dict[str, Any]' = <factory>, model_cv_kwargs: 'dict[str, Any]' = <factory>, model_key_override: 'str | None' = None, _fitted_model: 'Adapter | None' = None, _cv_results: 'list[CVResult] | None' = None, _diagnostics: 'CVDiagnostics | None' = None, _last_cv_kwargs: 'dict[str, Any] | None' = None, _last_fit_arg_ids: 'dict[str, int] | None' = None, factory_function: 'TModelFactory | None' = None)[source]

Bases: GeneratorSpec

factory_function: Callable[[...], Adapter] | None = None

fit(experiment: Experiment, data: Data | None = None, search_space: SearchSpace | None = None, optimization_config: OptimizationConfig | None = None, **model_kwargs: Any) → None[source]: Fits the specified model on the given experiment + data using the model kwargs set on the model spec, alongside any passed down as kwargs to this function (local kwargs take precedent)

model_enum: ModelRegistryBase | None = None

class ax.generation_strategy.model_spec.GeneratorSpec(model_enum: 'ModelRegistryBase', model_kwargs: 'dict[str, Any]' = <factory>, model_gen_kwargs: 'dict[str, Any]' = <factory>, model_cv_kwargs: 'dict[str, Any]' = <factory>, model_key_override: 'str | None' = None, _fitted_model: 'Adapter | None' = None, _cv_results: 'list[CVResult] | None' = None, _diagnostics: 'CVDiagnostics | None' = None, _last_cv_kwargs: 'dict[str, Any] | None' = None, _last_fit_arg_ids: 'dict[str, int] | None' = None)[source]

Bases: SortableBase, SerializationMixin

copy() → GeneratorSpec[source]: GeneratorSpec is both a spec and an object that performs actions. Copying is useful to avoid changes to a singleton model spec.

cross_validate(model_cv_kwargs: dict[str, Any] | None = None) → tuple[list[CVResult] | None, dict[str, dict[str, float]] | None][source]

Call cross_validate, compute_diagnostics and cache the results. If the model cannot be cross validated, warn and return None.

NOTE: If there are cached results, and the cache was computed using the same kwargs, this will return the cached results.

Parameters:: model_cv_kwargs – Optional kwargs to pass into cross_validate call. These are combined with self.model_cv_kwargs, with the model_cv_kwargs taking precedence over self.model_cv_kwargs.
Returns:: A tuple of CV results (observed vs predicted values) and the corresponding diagnostics.

property cv_results: list[CVResult] | None: Cached CV results from self.cross_validate() if it has been successfully called

property diagnostics: dict[str, dict[str, float]] | None: Cached CV diagnostics from self.cross_validate() if it has been successfully called

fit(experiment: Experiment, data: Data | None = None, **model_kwargs: Any) → None[source]: Fits the specified model on the given experiment + data using the model kwargs set on the model spec, alongside any passed down as kwargs to this function (local kwargs take precedent)

property fitted_model: Adapter: Returns the fitted Ax model, asserting fit() was called

property fixed_features: ObservationFeatures | None: Fixed generation features to pass into the Model’s .gen function.

gen(**model_gen_kwargs: Any) → GeneratorRun[source]

Generates candidates from the fitted model, using the model gen kwargs set on the model spec, alongside any passed as kwargs to this function (local kwargs take precedent)

NOTE: Model must have been fit prior to calling gen()

Parameters:

n – Integer representing how many arms should be in the generator run produced by this method. NOTE: Some underlying models may ignore the n and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that can differ from n.
pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.

model_cv_kwargs: dict[str, Any]

model_enum: ModelRegistryBase

model_gen_kwargs: dict[str, Any]

property model_key: str: Key string to identify the model used by this GeneratorSpec.

model_key_override: str | None = None

model_kwargs: dict[str, Any]

class ax.generation_strategy.model_spec.GeneratorSpecJSONEncoder(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, sort_keys=False, indent=None, separators=None, default=None)[source]

Bases: JSONEncoder

Generic encoder to avoid JSON errors in GeneratorSpec.__repr__

default(o: Any) → str[source]

Implement this method in a subclass such that it returns a serializable object for o, or calls the base implementation (to raise a TypeError).

For example, to support arbitrary iterators, you could implement default like this:

def default(self, o):
    try:
        iterable = iter(o)
    except TypeError:
        pass
    else:
        return list(iterable)
    # Let the base class default method raise the TypeError
    return super().default(o)

Model Selection

class ax.generation_strategy.best_model_selector.BestModelSelector[source]

Bases: ABC, Base

abstractmethod best_model(model_specs: list[GeneratorSpec]) → GeneratorSpec[source]

Return the best GeneratorSpec based on some criteria.

NOTE: The returned GeneratorSpec may be a different object than what was provided in the original list. It may be possible to clone and modify the original GeneratorSpec to produce one that performs better.

class ax.generation_strategy.best_model_selector.ReductionCriterion(*values)[source]

Bases: Enum

An enum for callables that are used for aggregating diagnostics over metrics and selecting the best diagnostic in SingleDiagnosticBestModelSelector.

NOTE: This is used to ensure serializability of the callables.

MAX: Callable[[ndarray | list[float] | list[ndarray]], ndarray[Any, dtype[_ScalarType_co]]] = functools.partial(<function max>)

MEAN: Callable[[ndarray | list[float] | list[ndarray]], ndarray[Any, dtype[_ScalarType_co]]] = functools.partial(<function mean>)

MIN: Callable[[ndarray | list[float] | list[ndarray]], ndarray[Any, dtype[_ScalarType_co]]] = functools.partial(<function min>)

class ax.generation_strategy.best_model_selector.SingleDiagnosticBestModelSelector(diagnostic: str, metric_aggregation: ReductionCriterion, criterion: ReductionCriterion, model_cv_kwargs: dict[str, Any] | None = None)[source]

Bases: BestModelSelector

Choose the best model using a single cross-validation diagnostic.

The input is a list of GeneratorSpec, each corresponding to one model. The specified diagnostic is extracted from each of the models, its values (each of which corresponds to a separate metric) are aggregated with the aggregation function, the best one is determined with the criterion, and the index of the best diagnostic result is returned.

Example

::

s = SingleDiagnosticBestModelSelector(: diagnostic=’Fisher exact test p’, metric_aggregation=ReductionCriterion.MEAN, criterion=ReductionCriterion.MIN, model_cv_kwargs={“untransform”: False},

) best_model = s.best_model(model_specs=model_specs)

Parameters:

diagnostic – The name of the diagnostic to use, which should be a key in CVDiagnostic.
metric_aggregation – ReductionCriterion applied to the values of the diagnostic for a single model to produce a single number.
criterion – ReductionCriterion used to determine which of the (aggregated) diagnostics is the best.
model_cv_kwargs – Optional dictionary of kwargs to pass in while computing the cross validation diagnostics.

best_model(model_specs: list[GeneratorSpec]) → GeneratorSpec[source]

Return the best GeneratorSpec based on the specified diagnostic.

Parameters:: model_specs – List of GeneratorSpec to choose from.
Returns:: The best GeneratorSpec based on the specified diagnostic.

Dispatch Utilities

ax.generation_strategy.dispatch_utils.calculate_num_initialization_trials(num_tunable_parameters: int, num_trials: int | None, use_batch_trials: bool) → int[source]

Applies rules from high to low priority

1 for batch trials.
At least 5
At most 1/5th of num_trials.
Twice the number of tunable parameters

ax.generation_strategy.dispatch_utils.choose_generation_strategy_legacy(search_space: SearchSpace, *, use_batch_trials: bool = False, enforce_sequential_optimization: bool = True, random_seed: int | None = None, torch_device: device | None = None, no_winsorization: bool = False, winsorization_config: None | WinsorizationConfig | dict[str, WinsorizationConfig] = None, derelativize_with_raw_status_quo: bool = False, no_bayesian_optimization: bool | None = None, force_random_search: bool = False, num_trials: int | None = None, num_initialization_trials: int | None = None, num_completed_initialization_trials: int = 0, max_initialization_trials: int | None = None, min_sobol_trials_observed: int | None = None, max_parallelism_cap: int | None = None, max_parallelism_override: int | None = None, optimization_config: OptimizationConfig | None = None, should_deduplicate: bool = False, use_saasbo: bool = False, verbose: bool | None = None, disable_progbar: bool | None = None, jit_compile: bool | None = None, experiment: Experiment | None = None, suggested_model_override: ModelRegistryBase | None = None, fit_out_of_design: bool = False) → GenerationStrategy[source]

Select an appropriate generation strategy based on the properties of the search space and expected settings of the experiment, such as number of arms per trial, optimization algorithm settings, expected number of trials in the experiment, etc.

Parameters:

search_space – SearchSpace, based on the properties of which to select the generation strategy.
use_batch_trials – Whether this generation strategy will be used to generate batched trials instead of 1-arm trials.
enforce_sequential_optimization – Whether to enforce that 1) the generation strategy needs to be updated with min_trials_observed observations for a given generation step before proceeding to the next one and 2) maximum number of trials running at once (max_parallelism) if enforced for the BayesOpt step. NOTE: max_parallelism_override and max_parallelism_cap settings will still take their effect on max parallelism even if enforce_sequential_optimization=False, so if those settings are specified, max parallelism will be enforced.
random_seed – Fixed random seed for the Sobol generator.
torch_device – The device to use for generation steps implemented in PyTorch (e.g. via BoTorch). Some generation steps (in particular EHVI-based ones for multi-objective optimization) can be sped up by running candidate generation on the GPU. If not specified, uses the default torch device (usually the CPU).
no_winsorization – Whether to apply the winsorization transform prior to applying other transforms for fitting the BoTorch model.
winsorization_config – Explicit winsorization settings, if winsorizing. Usually only upper_quantile_margin is set when minimizing, and only lower_quantile_margin when maximizing.
derelativize_with_raw_status_quo – Whether to derelativize using the raw status quo values in any transforms. This argument is primarily to allow automatic Winsorization when relative constraints are present. Note: automatic Winsorization will fail if this is set to False (or unset) and there are relative constraints present.
no_bayesian_optimization – Deprecated. Use force_random_search.
force_random_search – If True, quasi-random generation strategy will be used rather than Bayesian optimization.
num_trials – Total number of trials in the optimization, if known in advance.
num_initialization_trials – Specific number of initialization trials, if wanted. Typically, initialization trials are generated quasi-randomly.
max_initialization_trials – If num_initialization_trials unspecified, it will be determined automatically. This arg provides a cap on that automatically determined number.
num_completed_initialization_trials – The final calculated number of initialization trials is reduced by this number. This is useful when warm-starting an experiment, to specify what number of completed trials can be used to satisfy the initialization_trial requirement.
min_sobol_trials_observed – Minimum number of Sobol trials that must be observed before proceeding to the next generation step. Defaults to ceil(num_initialization_trials / 2).
max_parallelism_cap – Integer cap on parallelism in this generation strategy. If specified, max_parallelism setting in each generation step will be set to the minimum of the default setting for that step and the value of this cap. max_parallelism_cap is meant to just be a hard limit on parallelism (e.g. to avoid overloading machine(s) that evaluate the experiment trials). Specify only if not specifying max_parallelism_override.
max_parallelism_override – Integer, with which to override the default max parallelism setting for all steps in the generation strategy returned from this function. Each generation step has a max_parallelism value, which restricts how many trials can run simultaneously during a given generation step. By default, the parallelism setting is chosen as appropriate for the model in a given generation step. If max_parallelism_override is -1, no max parallelism will be enforced for any step of the generation strategy. Be aware that parallelism is limited to improve performance of Bayesian optimization, so only disable its limiting if necessary.
optimization_config – used to infer whether to use MOO and will be passed in to Winsorize via its transform_config in order to determine default winsorization behavior when necessary.
should_deduplicate – Whether to deduplicate the parameters of proposed arms against those of previous arms via rejection sampling. If this is True, the generation strategy will discard generator runs produced from the generation step that has should_deduplicate=True if they contain arms already present on the experiment and replace them with new generator runs. If no generator run with entirely unique arms could be produced in 5 attempts, a GenerationStrategyRepeatedPoints error will be raised, as we assume that the optimization converged when the model can no longer suggest unique arms.
use_saasbo – Whether to use SAAS prior for any GPEI generation steps.
verbose – Whether GP model should produce verbose logs. If not None, its value gets added to model_kwargs during generation_strategy construction. Defaults to True for SAASBO, else None. Verbose outputs are currently only available for SAASBO, so if verbose is not None for a different model type, it will be overridden to None with a warning.
disable_progbar – Whether GP model should produce a progress bar. If not None, its value gets added to model_kwargs during generation_strategy construction. Defaults to True for SAASBO, else None. Progress bars are currently only available for SAASBO, so if disable_probar is not None for a different model type, it will be overridden to None with a warning.
jit_compile – Whether to use jit compilation in Pyro when SAASBO is used.
experiment – If specified, _experiment attribute of the generation strategy will be set to this experiment (useful for associating a generation strategy with a given experiment before it’s first used to gen with that experiment). Can also provide optimization_config if it is not provided as an arg to this function.
suggested_model_override – If specified, this model will be used for the GP step and automatic selection will be skipped.
fit_out_of_design – Whether to include out-of-design points in the model.