ax.utils¶

Base¶

class ax.utils.common.base.Base[source]¶

Bases: object

Metaclass for core Ax classes. Provides an equality check and db_id property for SQA storage.

property db_id: Optional[int]¶

class ax.utils.common.base.SortableBase[source]¶

Bases: Base

Extension to the base class that also provides an inequality check.

Constants¶

class ax.utils.common.constants.Keys(value)[source]¶

Bases: str, Enum

Enum of reserved keys in options dicts etc, alphabetized.

NOTE: Useful for keys in dicts that correspond to kwargs to classes or functions and/or are used in multiple places.

ACQF_KWARGS = 'acquisition_function_kwargs'¶

AUTOSET_SURROGATE = 'autoset_surrogate'¶

BATCH_INIT_CONDITIONS = 'batch_initial_conditions'¶

CANDIDATE_SET = 'candidate_set'¶

CANDIDATE_SIZE = 'candidate_size'¶

COST_AWARE_UTILITY = 'cost_aware_utility'¶

COST_INTERCEPT = 'cost_intercept'¶

CURRENT_VALUE = 'current_value'¶

EXPAND = 'expand'¶

EXPECTED_ACQF_VAL = 'expected_acquisition_value'¶

FIDELITY_FEATURES = 'fidelity_features'¶

FIDELITY_WEIGHTS = 'fidelity_weights'¶

FRAC_RANDOM = 'frac_random'¶

FULL_PARAMETERIZATION = 'full_parameterization'¶

IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF = 'immutable_search_space_and_opt_config'¶

MAXIMIZE = 'maximize'¶

METADATA = 'metadata'¶

METRIC_NAMES = 'metric_names'¶

NUM_FANTASIES = 'num_fantasies'¶

NUM_INNER_RESTARTS = 'num_inner_restarts'¶

NUM_RESTARTS = 'num_restarts'¶

NUM_TRACE_OBSERVATIONS = 'num_trace_observations'¶

OBJECTIVE = 'objective'¶

ONLY_SURROGATE = 'only_surrogate'¶

OPTIMIZER_KWARGS = 'optimizer_kwargs'¶

PAIRWISE_PREFERENCE_QUERY = 'pairwise_pref_query'¶

PREFERENCE_DATA = 'preference_data'¶

PRIMARY_SURROGATE = 'primary'¶

PROJECT = 'project'¶

QMC = 'qmc'¶

RAW_INNER_SAMPLES = 'raw_inner_samples'¶

RAW_SAMPLES = 'raw_samples'¶

REFIT_ON_UPDATE = 'refit_on_update'¶

SAMPLER = 'sampler'¶

SEED_INNER = 'seed_inner'¶

SEQUENTIAL = 'sequential'¶

STATE_DICT = 'state_dict'¶

SUBCLASS = 'subclass'¶

SUBSET_MODEL = 'subset_model'¶

TASK_FEATURES = 'task_features'¶

TRIAL_COMPLETION_TIMESTAMP = 'trial_completion_timestamp'¶

WARM_START_REFITTING = 'warm_start_refitting'¶

X_BASELINE = 'X_baseline'¶

Decorator¶

class ax.utils.common.decorator.ClassDecorator[source]¶

Bases: ABC

Template for making a decorator work as a class level decorator. That decorator should extend ClassDecorator. It must implement __init__ and decorate_callable. See disable_logger.decorate_callable for an example. decorate_callable should call self._call_func() instead of directly calling func to handle static functions. Note: _call_func is still imperfect and unit tests should be used to ensure everything is working properly. There is a lot of complexity in detecting classmethods and staticmethods and removing the self argument in the right situations. For best results always use keyword args in the decorated class.

DECORATE_PRIVATE can be set to determine whether private methods should be decorated. In the case of a logging decorator, you may only want to decorate things the user calls. But in the case of a disable logging decorator, you may want to decorate everything to ensure no logs escape.

DECORATE_PRIVATE = True¶

abstract decorate_callable(func: Callable[[...], T]) → Callable[[...], T][source]¶

decorate_class(klass: T) → T[source]¶

Docutils¶

Support functions for sphinx et. al

ax.utils.common.docutils.copy_doc(src: Callable[[...], Any]) → Callable[[_T], _T][source]¶

A decorator that copies the docstring of another object

Since sphinx actually loads the python modules to grab the docstrings this works with both sphinx and the help function.

class Cat(Mamal):

  @property
  @copy_doc(Mamal.is_feline)
  def is_feline(self) -> true:
      ...

Equality¶

ax.utils.common.equality.dataframe_equals(df1: pandas.DataFrame, df2: pandas.DataFrame) → bool[source]¶

Compare equality of two pandas dataframes.

ax.utils.common.equality.datetime_equals(dt1: Optional[datetime], dt2: Optional[datetime]) → bool[source]¶

Compare equality of two datetimes, ignoring microseconds.

ax.utils.common.equality.equality_typechecker(eq_func: Callable) → Callable[source]¶

A decorator to wrap all __eq__ methods to ensure that the inputs are of the right type.

ax.utils.common.equality.object_attribute_dicts_equal(one_dict: Dict[str, Any], other_dict: Dict[str, Any], skip_db_id_check: bool = False) → bool[source]¶

Utility to check if all items in attribute dicts of two Ax objects are the same.

NOTE: Special-cases some Ax object attributes, like “_experiment” or “_model”, where full equality is hard to check.

Parameters:

• one_dict – First object’s attribute dict (obj.__dict__).

• other_dict – Second object’s attribute dict (obj.__dict__).

• skip_db_id_check – If True, will exclude the db_id attributes from the equality check. Useful for ensuring that all attributes of an object are equal except the ids, with which one or both of them are saved to the database (e.g. if confirming an object before it was saved, to the version reloaded from the DB).

ax.utils.common.equality.object_attribute_dicts_find_unequal_fields(one_dict: Dict[str, Any], other_dict: Dict[str, Any], fast_return: bool = True, skip_db_id_check: bool = False) → Tuple[Dict[str, Tuple[Any, Any]], Dict[str, Tuple[Any, Any]]][source]¶

Utility for finding out what attributes of two objects’ attribute dicts are unequal.

Parameters:

• one_dict – First object’s attribute dict (obj.__dict__).

• other_dict – Second object’s attribute dict (obj.__dict__).

• fast_return – Boolean representing whether to return as soon as a single unequal attribute was found or to iterate over all attributes and collect all unequal ones.

• skip_db_id_check – If True, will exclude the db_id attributes from the equality check. Useful for ensuring that all attributes of an object are equal except the ids, with which one or both of them are saved to the database (e.g. if confirming an object before it was saved, to the version reloaded from the DB).

Returns:

• attribute name to attribute values of unequal type (as a tuple),

• attribute name to attribute values of unequal value (as a tuple).

Return type:

Two dictionaries

ax.utils.common.equality.same_elements(list1: List[Any], list2: List[Any]) → bool[source]¶

Compare equality of two lists of core Ax objects.

Assumptions:

– The contents of each list are types that implement __eq__ – The lists do not contain duplicates

Checking equality is then the same as checking that the lists are the same length, and that one is a subset of the other.

Equality¶

ax.utils.common.executils.handle_exceptions_in_retries(no_retry_exceptions: Tuple[Type[Exception], ...], retry_exceptions: Tuple[Type[Exception], ...], suppress_errors: bool, check_message_contains: Optional[str], last_retry: bool, logger: Optional[Logger], wrap_error_message_in: Optional[str]) → Generator[None, None, None][source]¶

ax.utils.common.executils.retry_on_exception(exception_types: Optional[Tuple[Type[Exception], ...]] = None, no_retry_on_exception_types: Optional[Tuple[Type[Exception], ...]] = None, check_message_contains: Optional[List[str]] = None, retries: int = 3, suppress_all_errors: bool = False, logger: Optional[Logger] = None, default_return_on_suppression: Optional[Any] = None, wrap_error_message_in: Optional[str] = None, initial_wait_seconds: Optional[int] = None) → Optional[Any][source]¶

A decorator for instance methods or standalone functions that makes them retry on failure and allows to specify on which types of exceptions the function should and should not retry.

NOTE: If the argument suppress_all_errors is supplied and set to True, the error will be suppressed and default value returned.

Parameters:

• exception_types – A tuple of exception(s) types to catch in the decorated function. If none is provided, baseclass Exception will be used.

• no_retry_on_exception_types – Exception types to consider non-retryable even if their supertype appears in exception_types or the only exceptions to not retry on if no exception_types are specified.

• check_message_contains – A list of strings, against which to match error messages. If the error message contains any one of these strings, the exception will cause a retry. NOTE: This argument works in addition to exception_types; if those are specified, only the specified types of exceptions will be caught and retried on if they contain the strings provided as check_message_contains.

• retries – Number of retries to perform.

• suppress_all_errors – If true, after all the retries are exhausted, the error will still be suppressed and default_return_on_suppresion will be returned from the function. NOTE: If using this argument, the decorated function may not actually get fully executed, if it consistently raises an exception.

• logger – A handle for the logger to be used.

• default_return_on_suppression – If the error is suppressed after all the retries, then this default value will be returned from the function. Defaults to None.

• wrap_error_message_in – If raising the error message after all the retries, a string wrapper for the error message (useful for making error messages more user-friendly). NOTE: Format of resulting error will be: “<wrap_error_message_in>: <original_error_type>: <original_error_msg>”, with the stack trace of the original message.

• initial_wait_seconds – Initial length of time to wait between failures, doubled after each failure up to a maximum of 10 minutes. If unspecified then there is no wait between retries.

Kwargs¶

ax.utils.common.kwargs.consolidate_kwargs(kwargs_iterable: Iterable[Optional[Dict[str, Any]]], keywords: Iterable[str]) → Dict[str, Any][source]¶

Combine an iterable of kwargs into a single dict of kwargs, where kwargs by duplicate keys that appear later in the iterable get priority over the ones that appear earlier and only kwargs referenced in keywords will be used. This allows to combine somewhat redundant sets of kwargs, where a user-set kwarg, for instance, needs to override a default kwarg.

>>> consolidate_kwargs(
...     kwargs_iterable=[{'a': 1, 'b': 2}, {'b': 3, 'c': 4, 'd': 5}],
...     keywords=['a', 'b', 'd']
... )
{'a': 1, 'b': 3, 'd': 5}

ax.utils.common.kwargs.filter_kwargs(function: Callable, **kwargs: Any) → Any[source]¶

Filter out kwargs that are not applicable for a given function. Return a copy of given kwargs dict with only the required kwargs.

ax.utils.common.kwargs.get_function_argument_names(function: Callable, omit: Optional[List[str]] = None) → List[str][source]¶

Extract parameter names from function signature.

ax.utils.common.kwargs.get_function_default_arguments(function: Callable) → Dict[str, Any][source]¶

Extract default arguments from function signature.

ax.utils.common.kwargs.validate_kwarg_typing(typed_callables: List[Callable], **kwargs: Any) → None[source]¶

Check if keywords in kwargs exist in any of the typed_callables and if the type of each keyword value matches the type of corresponding arg in one of the callables

Note: this function expects the typed callables to have unique keywords for the arguments and will raise an error if repeat keywords are found.

ax.utils.common.kwargs.warn_on_kwargs(callable_with_kwargs: Callable, **kwargs: Any) → None[source]¶

Log a warning when a decoder function receives unexpected kwargs.

NOTE: This mainly caters to the use case where an older version of Ax is used to decode objects, serialized to JSON by a newer version of Ax (and therefore potentially containing new fields). In that case, the decoding function should not fail when encountering those additional fields, but rather just ignore them and log a warning using this function.

Logger¶

class ax.utils.common.logger.AxOutputNameFilter(name='')[source]¶

Bases: Filter

This is a filter which sets the record’s output_name, if not configured

filter(record: LogRecord) → bool[source]¶

Determine if the specified record is to be logged.

Returns True if the record should be logged, or False otherwise. If deemed appropriate, the record may be modified in-place.

ax.utils.common.logger.build_file_handler(filepath: str, level: int = 20) → StreamHandler[source]¶

Build a file handle that logs entries to the given file, using the same formatting as the stream handler.

Parameters:

• filepath – Location of the file to log output to. If the file exists, output will be appended. If it does not exist, a new file will be created.

• level – The log level. By default, sets level to INFO

Returns:

A logging.FileHandler instance

ax.utils.common.logger.build_stream_handler(level: int = 20) → StreamHandler[source]¶

Build the default stream handler used for most Ax logging. Sets default level to INFO, instead of WARNING.

Parameters:

level – The log level. By default, sets level to INFO

Returns:

A logging.StreamHandler instance

class ax.utils.common.logger.disable_logger(name: str, level: int = 40)[source]¶

Bases: ClassDecorator

decorate_callable(func: Callable[[...], T]) → Callable[[...], T][source]¶

class ax.utils.common.logger.disable_loggers(names: List[str], level: int = 40)[source]¶

Bases: ClassDecorator

decorate_callable(func: Callable[[...], T]) → Callable[[...], T][source]¶

ax.utils.common.logger.get_logger(name: str, level: int = 20, force_name: bool = False) → Logger[source]¶

Get an Axlogger.

To set a human-readable “output_name” that appears in logger outputs, add {“output_name”: “[MY_OUTPUT_NAME]”} to the logger’s contextual information. By default, we use the logger’s name

NOTE: To change the log level on particular outputs (e.g. STDERR logs), set the proper log level on the relevant handler, instead of the logger e.g. logger.handers[0].setLevel(INFO)

Parameters:

• name – The name of the logger.

• level – The level at which to actually log. Logs below this level of importance will be discarded

• force_name – If set to false and the module specified is not ultimately a descendent of the ax module specified by name, “ax.” will be prepended to name

Returns:

The logging.Logger object.

ax.utils.common.logger.make_indices_str(indices: Iterable[int]) → str[source]¶

Generate a string representation of an iterable of indices; if indices are contiguous, returns a string formatted like like ‘<min_idx> - <max_idx>’, otherwise a string formatted like ‘[idx_1, idx_2, …, idx_n’].

ax.utils.common.logger.set_stderr_log_level(level: int) → None[source]¶

Set the log level for stream handler, such that logs of given level are printed to STDERR by the root logger

Mock Torch¶

ax.utils.common.mock.mock_patch_method_original(mock_path: str, original_method: Callable[[...], T]) → MagicMock[source]¶

Context manager for patching a method returning type T on class C, to track calls to it while still executing the original method. There is not a native way to do this with mock.patch.

Result¶

class ax.utils.common.result.Err(value: E)[source]¶

Bases: Generic[T, E], Result[T, E]

Contains the error value.

property err: E¶

is_err() → bool[source]¶

is_ok() → bool[source]¶

map(op: Callable[[T], U]) → Result[U, E][source]¶

Maps a Result[T, E] to Result[U, E] by applying a function to a contained Ok value, leaving an Err value untouched. This function can be used to compose the results of two functions.

map_err(op: Callable[[E], F]) → Result[T, F][source]¶

Maps a Result[T, E] to Result[T, F] by applying a function to a contained Err value, leaving an Ok value untouched. This function can be used to pass through a successful result while handling an error.

map_or(default: U, op: Callable[[T], U]) → U[source]¶

Returns the provided default (if Err), or applies a function to the contained value (if Ok).

map_or_else(default_op: Callable[[], U], op: Callable[[T], U]) → U[source]¶

Maps a Result[T, E] to U by applying fallback function default to a contained Err value, or function op to a contained Ok value. This function can be used to unpack a successful result while handling an error.

property ok: None¶

unwrap() → NoReturn[source]¶

Returns the contained Ok value.

Because this function may raise an UnwrapError, its use is generally discouraged. Instead, prefer to handle the Err case explicitly, or call unwrap_or, unwrap_or_else, or unwrap_or_default.

unwrap_err() → E[source]¶

Returns the contained Err value.

Because this function may raise an UnwrapError, its use is generally discouraged. Instead, prefer to handle the Err case explicitly, or call unwrap_or, unwrap_or_else, or unwrap_or_default.

unwrap_or(default: T) → T[source]¶

Returns the contained Ok value or a provided default.

unwrap_or_else(op: Callable[[E], T]) → T[source]¶

Returns the contained Ok value or computes it from a Callable.

property value: E¶

class ax.utils.common.result.Ok(value: T)[source]¶

Bases: Generic[T, E], Result[T, E]

Contains the success value.

property err: None¶

is_err() → bool[source]¶

is_ok() → bool[source]¶

map(op: Callable[[T], U]) → Result[U, E][source]¶

Maps a Result[T, E] to Result[U, E] by applying a function to a contained Ok value, leaving an Err value untouched. This function can be used to compose the results of two functions.

map_err(op: Callable[[E], F]) → Result[T, F][source]¶

Maps a Result[T, E] to Result[T, F] by applying a function to a contained Err value, leaving an Ok value untouched. This function can be used to pass through a successful result while handling an error.

map_or(default: U, op: Callable[[T], U]) → U[source]¶

Returns the provided default (if Err), or applies a function to the contained value (if Ok).

map_or_else(default_op: Callable[[], U], op: Callable[[T], U]) → U[source]¶

Maps a Result[T, E] to U by applying fallback function default to a contained Err value, or function op to a contained Ok value. This function can be used to unpack a successful result while handling an error.

property ok: T¶

unwrap() → T[source]¶

Returns the contained Ok value.

Because this function may raise an UnwrapError, its use is generally discouraged. Instead, prefer to handle the Err case explicitly, or call unwrap_or, unwrap_or_else, or unwrap_or_default.

unwrap_err() → NoReturn[source]¶

Returns the contained Err value.

Because this function may raise an UnwrapError, its use is generally discouraged. Instead, prefer to handle the Err case explicitly, or call unwrap_or, unwrap_or_else, or unwrap_or_default.

unwrap_or(default: U) → T[source]¶

Returns the contained Ok value or a provided default.

unwrap_or_else(op: Callable[[E], T]) → T[source]¶

Returns the contained Ok value or computes it from a Callable.

property value: T¶

class ax.utils.common.result.Result[source]¶

Bases: Generic[T, E], ABC

A minimal implementation of a rusty Result monad. See https://doc.rust-lang.org/std/result/enum.Result.html for more information.

abstract property err: Optional[E]¶

abstract is_err() → bool[source]¶

abstract is_ok() → bool[source]¶

abstract map(op: Callable[[T], U]) → Result[U, E][source]¶

Maps a Result[T, E] to Result[U, E] by applying a function to a contained Ok value, leaving an Err value untouched. This function can be used to compose the results of two functions.

abstract map_err(op: Callable[[E], F]) → Result[T, F][source]¶

Maps a Result[T, E] to Result[T, F] by applying a function to a contained Err value, leaving an Ok value untouched. This function can be used to pass through a successful result while handling an error.

abstract map_or(default: U, op: Callable[[T], U]) → U[source]¶

Returns the provided default (if Err), or applies a function to the contained value (if Ok).

abstract map_or_else(default_op: Callable[[], U], op: Callable[[T], U]) → U[source]¶

Maps a Result[T, E] to U by applying fallback function default to a contained Err value, or function op to a contained Ok value. This function can be used to unpack a successful result while handling an error.

abstract property ok: Optional[T]¶

abstract unwrap() → T[source]¶

Returns the contained Ok value.

Because this function may raise an UnwrapError, its use is generally discouraged. Instead, prefer to handle the Err case explicitly, or call unwrap_or, unwrap_or_else, or unwrap_or_default.

abstract unwrap_err() → E[source]¶

Returns the contained Err value.

Because this function may raise an UnwrapError, its use is generally discouraged. Instead, prefer to handle the Err case explicitly, or call unwrap_or, unwrap_or_else, or unwrap_or_default.

abstract unwrap_or(default: T) → T[source]¶

Returns the contained Ok value or a provided default.

abstract unwrap_or_else(op: Callable[[E], T]) → T[source]¶

Returns the contained Ok value or computes it from a Callable.

abstract property value: Union[T, E]¶

exception ax.utils.common.result.UnwrapError[source]¶

Bases: Exception

Exception that indicates something has gone wrong in an unwrap call.

This should not happen in real world use and indicates a user has impropperly or unsafely used the Result abstraction.

Serialization¶

class ax.utils.common.serialization.SerializationMixin[source]¶

Bases: object

classmethod deserialize_init_args(args: Dict[str, Any], decoder_registry: Optional[Dict[str, Type]] = None, class_decoder_registry: Optional[Dict[str, Callable[[Dict[str, Any]], Any]]] = None) → Dict[str, Any][source]¶

Given a dictionary, deserialize the properties needed to initialize the object. Used for storage.

classmethod serialize_init_args(obj: SerializationMixin) → Dict[str, Any][source]¶

Serialize the properties needed to initialize the object. Used for storage.

ax.utils.common.serialization.callable_from_reference(path: str) → Callable[source]¶

Retrieves a callable by its path.

ax.utils.common.serialization.callable_to_reference(callable: Callable) → str[source]¶

Obtains path to the callable of form <module>.<name>.

ax.utils.common.serialization.extract_init_args(args: Dict[str, Any], class_: Type) → Dict[str, Any][source]¶

Given a dictionary, extract the arguments required for the given class’s constructor.

ax.utils.common.serialization.named_tuple_to_dict(data: Any) → Any[source]¶

Recursively convert NamedTuples to dictionaries.

ax.utils.common.serialization.serialize_init_args(obj: Any, exclude_fields: Optional[List[str]] = None) → Dict[str, Any][source]¶

Given an object, return a dictionary of the arguments that are needed by its constructor.

Testutils¶

Support functions for tests

class ax.utils.common.testutils.TestCase(methodName: str = 'runTest')[source]¶

Bases: TestCase

The base Ax test case, contains various helper functions to write unittests.

assertAxBaseEqual(first: Base, second: Base, msg: Optional[str] = None, skip_db_id_check: bool = False) → None[source]¶

Check that two Ax objects that subclass Base are equal or raise assertion error otherwise.

Parameters:

• first – Base-subclassing object to compare to second.

• second – Base-subclassing object to compare to first.

• msg – Message to put into the assertion error raised on inequality; if not specified, a default message is used.

• skip_db_id_check –

  If True, will exclude the db_id attributes from the equality check. Useful for ensuring that all attributes of an object are equal except the ids, with which one or both of them are saved to the database (e.g. if confirming an object before it was saved, to the

  version reloaded from the DB).

assertDictsAlmostEqual(a: Dict[str, Any], b: Dict[str, Any], consider_nans_equal: bool = False) → None[source]¶

Testing utility that checks that 1) the keys of a and b are identical, and that 2) the values of a and b are almost equal if they have a floating point type, considering NaNs as equal, and otherwise just equal.

Parameters:

• test – The test case object.

• a – A dictionary.

• b – Another dictionary.

• consider_nans_equal – Whether to consider NaNs equal when comparing floating point numbers.

assertEqual(first: Any, second: Any, msg: Optional[str] = None) → None[source]¶

Fail if the two objects are unequal as determined by the ‘==’ operator.

assertRaisesOn(exc: Type[Exception], line: Optional[str] = None, regex: Optional[str] = None) → ContextManager[None][source]¶

Assert that an exception is raised on a specific line.

setUp() → None[source]¶

Only show log messages of WARNING or higher while testing.

Ax prints a lot of INFO logs that are not relevant for unit tests.

static silence_stderr() → Generator[None, None, None][source]¶

A context manager that silences stderr for part of a test.

If any exception passes through this context manager the stderr will be printed, otherwise it will be discarded.

ax.utils.common.testutils.disable_profiler(f: Callable[[...], T]) → Callable[[...], T][source]¶

Not all tests are compatible with yappi. This wrapper can be used on a test to stop profiling instead of disabling it for the entire test case

Timeutils¶

ax.utils.common.timeutils.current_timestamp_in_millis() → int[source]¶

Grab current timestamp in milliseconds as an int.

ax.utils.common.timeutils.timestamps_in_range(start: datetime, end: datetime, delta: timedelta) → Generator[datetime, None, None][source]¶

Generator of timestamps in range [start, end], at intervals delta.

ax.utils.common.timeutils.to_ds(ts: datetime) → str[source]¶

Convert a datetime to a DS string.

ax.utils.common.timeutils.to_ts(ds: str) → datetime[source]¶

Convert a DS string to a datetime.

Typeutils¶

ax.utils.common.typeutils.checked_cast(typ: Type[T], val: V, exception: Optional[Exception] = None) → T[source]¶

Cast a value to a type (with a runtime safety check).

Returns the value unchanged and checks its type at runtime. This signals to the typechecker that the value has the designated type.

Like typing.cast check_cast performs no runtime conversion on its argument, but, unlike typing.cast, checked_cast will throw an error if the value is not of the expected type. The type passed as an argument should be a python class.

Parameters:

• typ – the type to cast to

• val – the value that we are casting

• exception – override exception to raise if typecheck fails

Returns:

the val argument, unchanged

ax.utils.common.typeutils.checked_cast_dict(key_typ: Type[K], value_typ: Type[V], d: Dict[X, Y]) → Dict[K, V][source]¶

Calls checked_cast on all keys and values in the dictionary.

ax.utils.common.typeutils.checked_cast_list(typ: Type[T], old_l: List[V]) → List[T][source]¶

Calls checked_cast on all items in a list.

ax.utils.common.typeutils.checked_cast_optional(typ: Type[T], val: Optional[V]) → Optional[T][source]¶

Calls checked_cast only if value is not None.

ax.utils.common.typeutils.checked_cast_to_tuple(typ: Tuple[Type[V], ...], val: V) → T[source]¶

Cast a value to a union of multiple types (with a runtime safety check). This function is similar to checked_cast, but allows for the type to be defined as a tuple of types, in which case the value is cast as a union of the types in the tuple.

Parameters:

• typ – the tuple of types to cast to

• val – the value that we are casting

Returns:

the val argument, unchanged

ax.utils.common.typeutils.not_none(val: Optional[T], message: Optional[str] = None) → T[source]¶

Unbox an optional type.

Parameters:

• val – the value to cast to a non None type

• message – optional override of the default error message

Returns:

val when val is not None

Return type:

V

Throws:

ValueError if val is None

ax.utils.common.typeutils.numpy_type_to_python_type(value: Any) → Any[source]¶

If value is a Numpy int or float, coerce to a Python int or float. This is necessary because some of our transforms return Numpy values.

ax.utils.common.typeutils.version_safe_check_type(argname: str, value: T, expected_type: Type[T]) → None[source]¶

Excecute the check_type function if it has the expected signature, otherwise warn. This is done to support newer versions of typeguard with minimal loss of functionality for users that have dependency conflicts

Typeutils Torch¶

ax.utils.common.typeutils_torch.torch_type_from_str(identifier: str, type_name: str) → Union[dtype, device, Size][source]¶

ax.utils.common.typeutils_torch.torch_type_to_str(value: Union[dtype, device, Size]) → str[source]¶

Converts torch types, commonly used in Ax, to string representations.

Docstring Checker¶

ax.utils.flake8_plugins.docstring_checker.A000(node: AST) → Error¶

class ax.utils.flake8_plugins.docstring_checker.DocstringChecker(tree, filename)[source]¶

Bases: object

A flake8 plug-in that makes sure all public functions have a docstring

fikename: str¶

name: str = 'docstring checker'¶

run()[source]¶

tree: Module¶

version: str = '1.0.0'¶

class ax.utils.flake8_plugins.docstring_checker.DocstringCheckerVisitor[source]¶

Bases: NodeVisitor

check_A000(node: AST) → None[source]¶

errors: List[Error]¶

visit_AsyncFunctionDef(node: ClassDef) → None[source]¶

visit_ClassDef(node: ClassDef) → None[source]¶

visit_FunctionDef(node: FunctionDef) → None[source]¶

class ax.utils.flake8_plugins.docstring_checker.Error(lineno, col, message, type)[source]¶

Bases: NamedTuple

col: int¶

Alias for field number 1

lineno: int¶

Alias for field number 0

message: str¶

Alias for field number 2

type: Type¶

Alias for field number 3

ax.utils.flake8_plugins.docstring_checker.is_copy_doc_call(c)[source]¶

Tries to guess if this is a call to the copy_doc decorator. This is a purely syntactic check so if the decorator was aliased as another name] or wrapped in another function we will fail.

ax.utils.flake8_plugins.docstring_checker.new_error(errorid: str, msg: str) → Callable[[AST], Error][source]¶

ax.utils.flake8_plugins.docstring_checker.should_check(filename)[source]¶

Synthetic Functions¶

class ax.utils.measurement.synthetic_functions.Aug_Branin[source]¶

Bases: SyntheticFunction

Augmented Branin function (3-dimensional with infinitely many global minima).

class ax.utils.measurement.synthetic_functions.Aug_Hartmann6[source]¶

Bases: Hartmann6

Augmented Hartmann6 function (7-dimensional with 1 global minimum).

class ax.utils.measurement.synthetic_functions.Branin[source]¶

Bases: SyntheticFunction

Branin function (2-dimensional with 3 global minima).

class ax.utils.measurement.synthetic_functions.FromBotorch(botorch_synthetic_function: SyntheticTestFunction)[source]¶

Bases: SyntheticFunction

property name: str¶

class ax.utils.measurement.synthetic_functions.Hartmann6[source]¶

Bases: SyntheticFunction

Hartmann6 function (6-dimensional with 1 global minimum).

class ax.utils.measurement.synthetic_functions.SyntheticFunction[source]¶

Bases: ABC

property domain: Any¶

f(X: ndarray) → Union[float, ndarray][source]¶

Synthetic function implementation.

Parameters:

X (numpy.ndarray) – an n by d array, where n represents the number of observations and d is the dimensionality of the inputs.

Returns:

an n-dimensional array.

Return type:

numpy.ndarray

property fmax: Any¶

property fmin: Any¶

property maximums: Any¶

property minimums: Any¶

property name: Any¶

property required_dimensionality: Any¶

ax.utils.measurement.synthetic_functions.from_botorch(botorch_synthetic_function: SyntheticTestFunction) → SyntheticFunction[source]¶

Utility to generate Ax synthetic functions from BoTorch synthetic functions.

ax.utils.measurement.synthetic_functions.informative_failure_on_none(func: Callable) → Any[source]¶

Plotting¶

ax.utils.notebook.plotting.init_notebook_plotting(offline: bool = False) → None[source]¶

Initialize plotting in notebooks, either in online or offline mode.

ax.utils.notebook.plotting.render(plot_config: AxPlotConfig, inject_helpers: bool = False) → None[source]¶

Render plot config.

Render¶

ax.utils.report.render.h2_html(text: str) → str[source]¶

Embed text in subheading tag.

ax.utils.report.render.h3_html(text: str) → str[source]¶

Embed text in subsubheading tag.

ax.utils.report.render.link_html(text: str, href: str) → str[source]¶

Embed text and reference address into link tag.

ax.utils.report.render.list_item_html(text: str) → str[source]¶

Embed text in list element tag.

ax.utils.report.render.p_html(text: str) → str[source]¶

Embed text in paragraph tag.

ax.utils.report.render.render_report_elements(experiment_name: str, html_elements: List[str], header: bool = True, offline: bool = False, notebook_env: bool = False) → str[source]¶

Generate Ax HTML report for a given experiment from HTML elements.

Uses Jinja2 for template. Injects Plotly JS for graph rendering.

Example:

html_elements = [
    h2_html("Subsection with plot"),
    p_html("This is an example paragraph."),
    plot_html(plot_fitted(gp_model, 'perf_metric')),
    h2_html("Subsection with table"),
    pandas_html(data.df),
]
html = render_report_elements('My experiment', html_elements)

Parameters:

• experiment_name – the name of the experiment to use for title.

• html_elements – list of HTML strings to render in report body.

• header – if True, render experiment title as a header. Meant to be used for standalone reports (e.g. via email), as opposed to served on the front-end.

• offline – if True, entire Plotly library is bundled with report.

• notebook_env – if True, caps the report width to 700px for viewing in a notebook environment.

Returns:

HTML string.

Return type:

str

ax.utils.report.render.table_cell_html(text: str, width: Optional[str] = None) → str[source]¶

Embed text or an HTML element into table cell tag.

ax.utils.report.render.table_heading_cell_html(text: str) → str[source]¶

Embed text or an HTML element into table heading cell tag.

ax.utils.report.render.table_html(table_rows: List[str]) → str[source]¶

Embed list of HTML elements into table tag.

ax.utils.report.render.table_row_html(table_cells: List[str]) → str[source]¶

Embed list of HTML elements into table row tag.

ax.utils.report.render.unordered_list_html(list_items: List[str]) → str[source]¶

Embed list of html elements into an unordered list tag.

Derivative GP¶

ax.utils.sensitivity.derivative_gp.get_KXX_inv(gp: Model) → Tensor[source]¶

Get the inverse matrix of K(X,X). :param gp: Botorch model.

Returns:

The inverse of K(X,X).

ax.utils.sensitivity.derivative_gp.get_KxX_dx(gp: Model, x: Tensor, kernel_type: str = 'rbf') → Tensor[source]¶

Computes the analytic derivative of the kernel K(x,X) w.r.t. x. :param gp: Botorch model. :param x: (n x D) Test points. :param kernel_type: Takes “rbf” or “matern”

Returns:

Tensor (n x D) The derivative of the kernel K(x,X) w.r.t. x.

ax.utils.sensitivity.derivative_gp.get_Kxx_dx2(gp: Model, kernel_type: str = 'rbf') → Tensor[source]¶

Computes the analytic second derivative of the kernel w.r.t. the training data :param gp: Botorch model. :param kernel_type: Takes “rbf” or “matern”

Returns:

Tensor (n x D x D) The second derivative of the kernel w.r.t. the training data.

ax.utils.sensitivity.derivative_gp.posterior_derivative(gp: Model, x: Tensor, kernel_type: str = 'rbf') → MultivariateNormal[source]¶

Computes the posterior of the derivative of the GP w.r.t. the given test points x. This follows the derivation used by GIBO in Sarah Muller, Alexander von Rohr, Sebastian Trimpe. “Local policy search with Bayesian optimization”, Advances in Neural Information Processing Systems 34, NeurIPS 2021. :param gp: Botorch model :param x: (n x D) Test points. :param kernel_type: Takes “rbf” or “matern”

Returns:

A Botorch Posterior.

Derivative Measures¶

class ax.utils.sensitivity.derivative_measures.GpDGSMGpMean(model: Model, bounds: Tensor, derivative_gp: bool = False, kernel_type: Optional[str] = None, Y_scale: float = 1.0, num_mc_samples: int = 10000, input_qmc: bool = False, dtype: dtype = torch.float64, num_bootstrap_samples: int = 1)[source]¶

Bases: object

aggregation(transform_fun: Callable[[Tensor], Tensor]) → Tensor[source]¶

bootstrap_indices: Optional[Tensor] = None¶

gradient_absolute_measure() → Tensor[source]¶

Computes the gradient absolute measure:

Returns:

if self.num_bootstrap_samples > 1

Tensor: (values, var_mc, stderr_mc) x dim

else

Tensor: (values) x dim

gradient_measure() → Tensor[source]¶

Computes the gradient measure:

Returns:

if self.num_bootstrap_samples > 1

Tensor: (values, var_mc, stderr_mc) x dim

else

Tensor: (values) x dim

gradients_square_measure() → Tensor[source]¶

Computes the gradient square measure:

Returns:

if num_bootstrap_samples > 1

Tensor: (values, var_mc, stderr_mc) x dim

else

Tensor: (values) x dim

mean_gradients: Optional[Tensor] = None¶

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

class ax.utils.sensitivity.derivative_measures.GpDGSMGpSampling(model: Model, bounds: Tensor, num_gp_samples: int, derivative_gp: bool = False, kernel_type: Optional[str] = None, Y_scale: float = 1.0, num_mc_samples: int = 10000, input_qmc: bool = False, gp_sample_qmc: bool = False, dtype: dtype = torch.float64, num_bootstrap_samples: int = 1)[source]¶

Bases: GpDGSMGpMean

aggregation(transform_fun: Callable[[Tensor], Tensor]) → Tensor[source]¶

samples_gradients: Optional[Tensor] = None¶

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

ax.utils.sensitivity.derivative_measures.compute_derivatives_from_model_list(model_list: List[Model], bounds: Tensor, **kwargs: Any) → Tensor[source]¶

Computes average derivatives of a list of models on a bounded domain. Estimation is according to the GP posterior mean function.

Parameters:

• model_list – A list of m botorch.models.model.Model types for which to compute the average derivative.

• bounds – A 2 x d Tensor of lower and upper bounds of the domain of the models.

• kwargs – Passed along to GpDGSMGpMean.

Returns:

A (m x d) tensor of gradient measures.

Sobol Measures¶

ax.utils.sensitivity.sobol_measures.GaussianLinkMean(mean: Tensor, var: Tensor) → Tensor[source]¶

ax.utils.sensitivity.sobol_measures.ProbitLinkMean(mean: Tensor, var: Tensor) → Tensor[source]¶

class ax.utils.sensitivity.sobol_measures.SobolSensitivity(bounds: Tensor, input_function: Optional[Callable[[Tensor], Tensor]] = None, num_mc_samples: int = 10000, input_qmc: bool = False, second_order: bool = False, num_bootstrap_samples: int = 1, bootstrap_array: bool = False)[source]¶

Bases: object

evalute_function(f_A_B_ABi: Optional[Tensor] = None) → None[source]¶

evaluates the objective function and devides the evaluation into

torch.Tensors needed for the indices computation.

Parameters:

f_A_B_ABi – Function evaluations on the entire grid of size M(d+2).

first_order_indices() → Tensor[source]¶

Computes the first order Sobol indices:

Returns:

if num_bootstrap_samples>1

Tensor: (values,var_mc,stderr_mc)x dim

else

Tensor: (values)x dim

generate_all_input_matrix() → Tensor[source]¶

second_order_indices(first_order_idxs: Optional[Tensor] = None, first_order_idxs_btsp: Optional[Tensor] = None) → Tensor[source]¶

Computes the Second order Sobol indices: :param first_order_idxs: Tensor of first order indices. :param first_order_idxs_btsp: Tensor of all first order indices given by bootstrap.

Returns:

if num_bootstrap_samples>1

Tensor: (values,var_mc,stderr_mc)x dim

else

Tensor: (values)x dim

total_order_indices() → Tensor[source]¶

Computes the total Sobol indices:

Returns:

if num_bootstrap_samples>1

Tensor: (values,var_mc,stderr_mc)x dim

else

Tensor: (values)x dim