Search spaces

Search spaces#

Search spaces are JAX PyTree leaves and may be nested in dicts, lists, tuples, or other PyTree containers. Candidate batches preserve the same structure and add a leading n_parallel axis to every leaf.

import jax.numpy as jnp

from hyperoptax import (
    DiscreteSpace,
    LinearSpace,
    LogSpace,
    QLinearSpace,
    QLogSpace,
)

space = {
    "optimizer": {
        "learning_rate": LogSpace(1e-5, 1e-1),
        "weight_decay": LinearSpace(0.0, 0.1),
    },
    "depth": QLinearSpace(2, 8, datatype=jnp.int32),
    "width": QLogSpace(32, 1024, datatype=jnp.int32),
    "batch_size": DiscreteSpace((32, 64, 128, 256)),
}

DiscreteSpace values must be numeric. Encode string categories as numbers and map them back to labels outside the vmapped objective when needed.

API#

class hyperoptax.spaces.Space[source]#

Bases: ABC

Abstract base class for hyperparameter search spaces.

class hyperoptax.spaces.LinearSpace(lower_bound, upper_bound)[source]#

Bases: Space

Uniform continuous space over [lower_bound, upper_bound].

Parameters:
lower_bound#

Inclusive lower bound of the interval.

Type:

float

upper_bound#

Exclusive upper bound of the interval.

Type:

float

class hyperoptax.spaces.DiscreteSpace(values)[source]#

Bases: Space

Discrete space over a fixed set of values.

Samples uniformly from values. transform snaps any continuous value to the nearest element, which is useful when discrete candidates are generated via continuous optimization (e.g. in BayesianSearch).

Parameters:

values (tuple)

values#

Tuple of candidate values to sample from.

Type:

tuple

class hyperoptax.spaces.LogSpace(lower_bound, upper_bound, base=10)[source]#

Bases: LinearSpace

Log-uniform continuous space over [lower_bound, upper_bound].

Samples uniformly in log space so that each order of magnitude receives equal probability mass. Useful for learning rates and other scale parameters that span several orders of magnitude.

Parameters:
lower_bound#

Inclusive lower bound (in original scale, e.g. 1e-5).

Type:

float

upper_bound#

Exclusive upper bound (in original scale, e.g. 1e-1).

Type:

float

base#

Logarithm base (default 10). Must be greater than 1.

Type:

float

class hyperoptax.spaces.QLinearSpace(lower_bound, upper_bound, datatype=<class 'jax.numpy.int32'>)[source]#

Bases: LinearSpace

Quantized (integer) variant of LinearSpace.

Samples uniformly from [lower_bound, upper_bound] and rounds to the nearest integer. Use this for discrete integer hyperparameters with a uniform prior (e.g. number of layers, batch size).

Parameters:
lower_bound#

Inclusive lower bound.

Type:

float

upper_bound#

Exclusive upper bound.

Type:

float

datatype#

Integer dtype used after rounding (default jnp.int32).

Type:

type

class hyperoptax.spaces.QLogSpace(lower_bound, upper_bound, base=10, datatype=<class 'jax.numpy.int32'>)[source]#

Bases: LogSpace

Quantized (integer) variant of LogSpace.

Samples in log space and rounds to the nearest integer. Use this for integer hyperparameters whose scale spans orders of magnitude (e.g. number of hidden units, number of warmup steps).

Parameters:
Contents