Write a sampler

A sampler answers one question — what should the next trial try? — and it is a public trait, not an internal detail: everything a built-in sampler is given, an out-of-tree sampler is given too. This page goes from nothing to a working, registered sampler.

The worked example is a real one. examples/rust/custom_sampler.rs implements stratified uniform search, registers it on a study, runs 64 trials under it, and then asserts the property it was written for. It is executed on every commit by cargo dev run-examples, so nothing on this page is a sketch. Run it yourself with cargo run -p atune --example custom_sampler.

One tab, not two — implementing the trait is a Rust exercise

The Python bindings hand out sampler handles: atune.samplers.Tpe() returns an opaque, frozen object wrapping an Arc<dyn Sampler> that was built in Rust, and there is no route from a Python class back to the trait. So a custom sampler is written in Rust and then used from either language. That is why this page's blocks have a single tab, and why examples/parity.toml records this example as unpaired rather than pretending a Python arm is coming.

What the trait asks for

Eight methods, of which four you must write and four have defaults that are already correct for a stateless sampler:

Method	Required	What it does
infer_relative_space	yes	Names the joint space to sample in one shot. Usually the study's declared schema; a define-by-run study infers it from history, which is why it can fail. An empty schema is legal and means sample everything independently.
sample_relative	yes	Draws that whole space at once — the seam where a sampler models correlations between parameters. Must bind every parameter in the space and nothing else.
sample_independent	yes	Draws one parameter the relative space did not cover: the fallback for define-by-run code that suggests a name nobody has seen.
reseed	yes	Reseeds auxiliary randomness so two workers do not draw the same numbers. Per-trial seeds arrive through the trial instead, so for a stateless sampler the correct body is empty — and every built-in sampler's body is empty today, because nothing in the study loop calls it. Do not build a design that needs it.
after_trial	no	Called once a trial reaches a terminal state. Where a history-dependent sampler updates its model. An error here does not fail the trial.
state	no	The sampler's persistable state, or None. The framework writes it to storage so another process can pick the study up.
snapshots_space	no	Answers whether the sampler takes a one-time snapshot of the space at construction, as Grid does. The default false is correct unless yours enumerates a fixed set up front — a study that declares an open search space refuses a snapshotting sampler at create, exactly as it refuses one whose state() is Some (a model in fixed-box coordinates cannot follow a growing bound).
restore_state	no	Called once, before any sampling, when a study is created or resumed.

The trait is object-safe and Send + Sync: a study holds one Arc<dyn Sampler> and calls it from every worker thread. Any interior mutability therefore goes behind a lock — and a sampler that needs none is a better sampler.

What to import, since the excerpts below begin after the example's own use block. Against atune_core — the normal dependency for a plugin — atune_core::prelude::* covers the trait and every type in its signatures, and the two names it does not carry are atune_core::seed::seed_for_name (the per-name mixing, deliberately one path segment away) and atune_core::study::Study, which you need to run a study in a test of your sampler. Against the facade, as the example does, atune::prelude::* carries Study, StudyView and the built-ins, and Assignment and TrialMeta come from atune::{space, trial}.

The implementation

Stratified splits each parameter's support into k equal strata and sends trial n into stratum n mod k, jittered inside it. Sixteen uniform draws can easily leave a third of a range unvisited; sixteen draws in eight strata cannot.

Rust

/// Stratified uniform sampling: trial *n* draws from stratum *n* mod `strata`.
///
/// **Stateless**, which is the strongest determinism promise the framework
/// offers: the assignment for trial number *n* is a pure function of
/// `(study seed, n, space)`, so a 16-worker study evaluates exactly the
/// configuration set a single-worker one does. Nothing is persisted
/// ([`state`](Sampler::state) stays `None`) and there is nothing to restore.
#[derive(Debug, Clone, Copy)]
struct Stratified {
    /// How many equal parts each parameter's support is split into.
    strata: u32,
}

impl Stratified {
    /// Creates the sampler.
    ///
    /// # Errors
    ///
    /// [`Error::InvalidSpace`] if `strata` is zero: "no strata" is not a
    /// degenerate stratification, it is a division by zero, and a constructor
    /// that accepted it would fail later inside a worker thread.
    fn new(strata: u32) -> Result<Self> {
        if strata == 0 {
            return Err(Error::InvalidSpace(
                "a stratified sampler needs at least one stratum".to_owned(),
            ));
        }
        Ok(Self { strata })
    }

    /// Draws one parameter — the whole sampler in one function.
    ///
    /// `trial.number` picks the stratum and `seed_for_name(trial.sampler_seed,
    /// name)` jitters inside it, so two parameters of one trial land in the same
    /// stratum index but at different points of it, and the same parameter of two
    /// consecutive trials lands in different strata.
    ///
    /// # Errors
    ///
    /// Whatever [`Distribution::from_unit`] reports for a malformed
    /// distribution, propagated verbatim.
    fn draw(self, trial: &TrialMeta, name: &str, dist: &Distribution) -> Result<ParamValue> {
        let stratum = trial.number.get() % u64::from(self.strata);
        let jitter = fraction(seed_for_name(trial.sampler_seed, name));
        #[allow(
            clippy::cast_precision_loss,
            reason = "`stratum` is smaller than `strata`, a u32, so the conversion is exact"
        )]
        let unit = (stratum as f64 + jitter) / f64::from(self.strata);
        dist.from_unit(unit)
    }
}

/// A seed's top 53 bits as a fraction in `[0, 1)`.
///
/// 53 bits is what an `f64` holds exactly, so the division by 2^53 is the only
/// arithmetic here and it cannot round.
fn fraction(seed: u64) -> f64 {
    #[allow(
        clippy::cast_precision_loss,
        reason = "`seed >> 11` is 53 bits, which an f64 holds exactly — that is the point of the shift"
    )]
    let bits = (seed >> 11) as f64;
    bits / TWO_POW_53
}

impl Sampler for Stratified {
    /// The study's declared space, or an empty one.
    ///
    /// An empty schema is legal and means "sample everything independently": a
    /// define-by-run study that declares no space sends every parameter through
    /// [`sample_independent`](Sampler::sample_independent) instead, and this
    /// sampler answers both the same way.
    fn infer_relative_space(&self, study: &StudyView) -> Result<SpaceSchema> {
        Ok(study.space().cloned().unwrap_or_else(SpaceSchema::empty))
    }

    /// Draws every parameter of the joint space.
    ///
    /// The whole space arrives at once so that a sampler *can* model correlations
    /// between parameters. This one does not — it is stratified, not joint — but
    /// the seam is where a sampler that does would put it.
    fn sample_relative(
        &self,
        _study: &StudyView,
        trial: &TrialMeta,
        space: &SpaceSchema,
    ) -> Result<Assignment> {
        let mut params = Assignment::new();
        for spec in space {
            params.insert(
                spec.name(),
                self.draw(trial, spec.name(), spec.distribution())?,
            );
        }
        Ok(params)
    }

    /// Draws one parameter the relative space did not cover.
    ///
    /// Identical to what [`sample_relative`](Sampler::sample_relative) would
    /// have produced for it — deliberately: a parameter that migrates between the
    /// declared space and the define-by-run fallback must keep its value.
    fn sample_independent(
        &self,
        _study: &StudyView,
        trial: &TrialMeta,
        name: &str,
        dist: &Distribution,
    ) -> Result<ParamValue> {
        self.draw(trial, name, dist)
    }

    /// Nothing to reseed.
    ///
    /// This hook exists for a sampler with auxiliary randomness of its own, so
    /// that two workers do not draw the same numbers. A stateless sampler takes
    /// every seed it uses from [`TrialMeta`], which the framework has already
    /// made worker-independent, so the correct implementation is empty — not a
    /// hidden `self.rng = …` that would make the draw depend on which worker
    /// called it.
    fn reseed(&self, _seed: u64) {}
}

Three things in there are the trait's real requirements rather than this sampler's taste, and getting any of them wrong produces a sampler that works and is not reproducible:

1. Derive determinism, never store it. Each trial arrives with a sampler_seed that is already a pure function of (study seed, trial number). Mixing it with the parameter name through seed_for_name — the derivation Random uses — makes a value depend on (study seed, trial number, parameter name) and on nothing else: not on the worker that ran the trial, not on the order the space is iterated in, not on which other parameters exist. Those are the only two randomness inputs you get: the trial's own seed and the parameter's name, plus the trial number as an index. The study view carries the directions, the declared space and the trials — not the study seed — so a draw that has to be correlated with another trial's draw can only get there by reading that trial's recorded value back out of the view, which costs you the stateless row of the table below. Know which you are building before you write it.
2. Do not draw sequentially. A sampler that pulls its values from one per-trial stream changes every parameter's value when the search space gains, loses or reorders a parameter, silently invalidating a study's history. Per-name derivation is what buys stability under a space edit.
3. Sample the unit interval, not the parameter. All the awkwardness — log scaling, step snapping, integer rounding, categorical labels — lives in Distribution::from_unit, so a sampler produces an f64 in [0, 1) and lets the distribution do the rest. It is also why this sampler works unchanged for an int, a bool or a categorical parameter, without a single branch on kind.

Registering it

One line. The builder takes an Arc<dyn Sampler>, and your type goes exactly where Tpe or Random would — there is no registry to add to, no attribute to apply, and nothing to rebuild inside atune:

Rust

// Registering a sampler is one line: the study takes an `Arc<dyn Sampler>`,
// and an out-of-tree type goes in exactly where `Tpe` or `Random` would.
let study = Study::builder()
    .sampler(Arc::new(Stratified::new(STRATA)?))
    .budget(Budget::trials(TRIALS))
    .create(
        StudyConfig::new(STUDY)
            .with_seed(SEED)
            .with_space(space.clone()),
    )?;

study.optimize(|ctx| {
    let x = ctx.suggest_f64("x", LOW..=HIGH, Scale::Linear)?;
    let y = ctx.suggest_f64("y", LOW..=HIGH, Scale::Linear)?;
    Ok(bowl(x, y).into())
})?;

Proving it

A sampler that runs is not a sampler that works. This one reads the stratum back out of the stored parameters — to_unit is from_unit's inverse — so the check tests what the study recorded rather than what the sampler believes:

Rust

// The property the sampler exists for, checked rather than claimed: trial `n`
// drew its `x` from stratum `n mod STRATA`, and every stratum was visited the
// same number of times. `to_unit` is `from_unit`'s inverse, so this reads the
// stratum back out of the *stored* parameter rather than out of the sampler.
let view = study.view()?;
let distribution = space
    .distribution("x")
    .ok_or_else(|| Error::NotFound("the space has no `x`".to_owned()))?;
let mut visits: BTreeMap<u64, u64> = BTreeMap::new();
for trial in view.completed() {
    let value = trial
        .params
        .get("x")
        .ok_or_else(|| Error::NotFound("a trial has no `x`".to_owned()))?;
    let unit = distribution.to_unit(value)?;
    let stratum = trial.number.get() % u64::from(STRATA);
    let (low, high) = stratum_bounds(stratum);
    if unit < low || unit >= high {
        return Err(Error::Conflict(format!(
            "trial {} drew `x` at unit {unit:?}, outside stratum {stratum} [{low:?}, {high:?})",
            trial.number.get()
        )));
    }
    *visits.entry(stratum).or_default() += 1;
}

let expected = TRIALS / u64::from(STRATA);
let counts: Vec<u64> = visits.values().copied().collect();
println!("{TRIALS} trials, {STRATA} strata, seed {SEED}");
println!("visits per stratum of `x`: {counts:?} (expected {expected} each)");
for stratum in 0..u64::from(STRATA) {
    let count = visits.get(&stratum).copied().unwrap_or_default();
    if count != expected {
        return Err(Error::Conflict(format!(
            "stratum {stratum} was visited {count} times, not {expected}"
        )));
    }
}

That is the shape worth copying: state the property your sampler exists for, then assert it from storage. Coverage, spacing, population invariants and "every configuration exactly once" are all reachable this way, and none of them needs access to atune's internals.

Say which determinism grade you are

Every sampler in this repository states its own standing, and yours should too, because a user cannot infer it from the outside:

Grade	The promise	Yours if
Stateless	The assignment for trial n is a pure function of (study seed, n, space). A 16-worker study evaluates exactly the configuration set a 1-worker study does.	You read only the trial's metadata and the space. Stratified, Random, Grid and Qmc are here.
History-dependent	Single worker, fixed seed: reproducible. Several workers: replayable — the recorded history explains every draw, but a rerun may differ, because which trials had finished when you were asked is thread scheduling.	You read the study view. Tpe is here.
Stateful	The same as history-dependent, plus: a study torn down and resumed must reproduce, bit for bit, what a never-closed study would have sampled next.	You carry something that cannot be recomputed from history. Dehb, Nsga2 and Cmaes are here.

The boundary is not a formality — Determinism is the contract, and the difference between the first row and the other two is whether a user may parallelise your sampler and still expect the same study.

If you keep state

Do not keep it where the framework cannot see it. Return it from state as a typed, versioned blob and take it back in restore_state; the framework persists it under the study's sampler scope and hands it back on resume. The two must round-trip exactly, and a blob whose version you no longer understand is an error to report, not a struct to guess at.

That is the whole reason state is data here rather than object innards: a pickled sampler cannot cross a process boundary, and a study that resumes with a half-restored population is worse than one that refuses to resume.

Where your sampler is allowed to live

atune_core is wasm32-clean, and a gate proves it in every feature combination. A sampler that reaches for the wall clock, the filesystem, threads or a heavy numerical dependency therefore does not belong in the core — that is exactly why cmaes, gp and carbs are facade features rather than core modules.

For your own crate the rule reads: depend on atune_core if you need only the traits, keep anything OS-shaped or dependency-heavy behind a cargo feature of your own, and your sampler stays usable in a wasm build. Publish a plugin crate is the packaging half of this.

Errors, and what they mean to the loop

Sampling is fallible on purpose: an infallible draw would force implementations to panic when a space is exhausted, and a panicking sampler takes the study down with it. Return Error::Sampler for a failure to draw, Error::SpaceExhausted when an exhaustive sampler has handed out every configuration — a normal end to a study, not a fault — and Error::InvalidSpace from your constructor for a configuration that can never work, so it fails where the user wrote it rather than inside a worker thread.

Where to go next

If you want to	Go to
Understand the seam before implementing it	Samplers and schedulers
Know exactly what determinism promises	Determinism
Decide a running trial's fate instead of its parameters	Write a scheduler
Ship it as a crate someone can depend on	Publish a plugin crate
Read the trait, method by method	Rust API