Open in GitHubRun in Google Colab

Quickstart

Ax is an open-source platform for adaptive experimentation, a technique used to efficiently tune parameters in complex systems. This guide will walk through installation, core concepts, and basic usage of Ax.

Installation

We recommend using pip to install Ax.

pip install ax-platform

Core concepts

• Experiment: A process of iteratively suggesting and evaluating parameters to improve some objective.
• Parameter: A variable that can be adjusted -- a collection of these form the space we are searching over during the optimization.
• Objective: The value being optimized.
• Trial: A set of parameters and its associated objectrive.
• Client: An object that manages the experiment and provides methods for interacting with it.

from ax import Client, RangeParameterConfig

# 1. Initialize the Client.
client = Client()

# 2. Configure where Ax will search.
client.configure_experiment(
    name="booth_function",
    parameters=[
        RangeParameterConfig(
            name="x1",
            bounds=(-10.0, 10.0),
            parameter_type="float",
        ),
        RangeParameterConfig(
            name="x2",
            bounds=(-10.0, 10.0),
            parameter_type="float",
        ),
    ],
)

# 3. Configure a metric Ax will target (see other Tutorials for adding constraints,
# multiple objectives, tracking metrics etc.)
client.configure_optimization(objective="-1 * booth")

# 4 Conduct the experiment with 20 trials: get each trial from Ax, evaluate the
# objective function, log data back to Ax.
for _ in range(20):
    # Use higher value of `max_trials` to run trials in parallel.
    for trial_index, parameters in client.get_next_trials(max_trials=1).items():
        client.complete_trial(
            trial_index=trial_index,
            raw_data={
                "booth": (parameters["x1"] + 2 * parameters["x2"] - 7) ** 2
                + (2 * parameters["x1"] + parameters["x2"] - 5) ** 2
            },
        )

# 5. Obtain the best-performing configuration; the true minimum for the booth
# function is at (1, 3)
client.get_best_parameterization()

Output:

/opt/hostedtoolcache/Python/3.12.10/x64/lib/python3.12/site-packages/pyro/ops/stats.py:527: SyntaxWarning: invalid escape sequence 'g'
we have :math:ES^{*}(P,Q) ge ES^{*}(Q,Q) with equality holding if and only if :math:P=Q, i.e.
({'x1': 1.2647782039987052, 'x2': 2.2390498140801878},
{'booth': (np.float64(0.3562669939345824), np.float64(6.583455389751908))},
18,
'18_0')