NRPT

The primary interface is autotuning: autotune / autosample discover the chain count, local-exploration count, and schedule together, then draw from the target. The building-block tuners and the core single run are below them.

Autotuning

hamon.autosample(key: jax.Array, *, n_samples: int, n_warmup: int = 0, steps_per_sample: int = 1, **autotune_kwargs) -> tuple[jax.Array, AutotuneReport]

One-shot: :func:autotune then draw — returns (samples, report).

The convenience entry for "give me samples." Forwards all keyword arguments to :func:autotune (ebm/program or factories, init_factory, beta_range, device, …), then draws n_samples from the tuned plan. For repeated draws from one tuned configuration, call :func:autotune once and reuse :meth:NRPTPlan.sample.

hamon.autotune

Full NRPT autotuning: orchestrate chain count, exploration, and schedule.

The one-call front door (autotune / autosample) and its result objects, pure host orchestration over the tuners in hamon.tuning and the jitted core in hamon.nrpt.

DeviceLike module-attribute

Represent a PEP 604 union type

E.g. for int | str

__cached__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__doc__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__file__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__name__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__package__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

annotations module-attribute

AutotuneReport

Diagnostics from an :func:autotune run.

Attributes:

Name	Type	Description
n_chains		discovered chain count N.
gibbs_steps_per_round		discovered local-exploration count n_expl.
Lambda		estimated global communication barrier Λ.
betas		the final tuned β ladder.
device		the resolved device (string) or None.
chain_history		per-probe records from the N search (:func:tune_chains).
exploration		the :func:tune_exploration result dict, or None when the n_expl search was skipped.
round_trip_diagnostics		round-trip summary from the final production run.
total_round_trips		total completed round trips observed during the final production run (summed across chains), or None.
production_rounds		number of rounds the production run used (the window total_round_trips and tau_observed were measured over).

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__dict__ class-attribute

Read-only proxy of a mapping.

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__weakref__ property

list of weak references to the object

summary() -> str

Human-readable multi-line summary.

__eq__(other)

__init__(n_chains: int, gibbs_steps_per_round: int, Lambda: float, betas: np.ndarray, device: str | None, chain_history: list, exploration: dict | None, round_trip_diagnostics: dict | None, total_round_trips: int | None = None, production_rounds: int | None = None) -> None

__repr__()

NRPTPlan

A tuned NRPT configuration plus a warm cold-chain state.

Returned by :func:autotune. Holds the discovered hyperparameters (N, schedule, n_expl) and an equilibrated cold-chain state, so :meth:sample can draw repeatedly and cheaply — no re-tuning, reusing the compiled loop.

Attributes:

Name	Type	Description
n_chains	/ betas / gibbs_steps_per_round / Lambda	the tuned config.
device	/ betas / gibbs_steps_per_round / Lambda	the resolved device (or None).
report	/ betas / gibbs_steps_per_round / Lambda	the :class:AutotuneReport.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__dict__ class-attribute

Read-only proxy of a mapping.

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__weakref__ property

list of weak references to the object

sample(key: jax.Array, n_samples: int, *, n_warmup: int = 0, steps_per_sample: int = 1) -> jax.Array

Draw n_samples from the target (cold chain) — cheap and repeatable.

Runs single-chain block Gibbs at the tuned cold β from the stored warm state (the established post-NRPT draw). Returns a (n_samples, n_nodes) array; call again with a fresh key for more, with no re-tuning.

__eq__(other)

__init__(n_chains: int, betas: np.ndarray, gibbs_steps_per_round: int, Lambda: float, device: Any, report: AutotuneReport, _cold_program: BlockSamplingProgram, _warm_state: list, _clamp_state: list, _obs_block: Any) -> None

__repr__()

AbstractEBM

Something that has a well-defined energy function (map from a state to a scalar).

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__init__() -> None

Initialize self. See help(type(self)) for accurate signature.

energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], _State]], blocks: BlockSpec | list[Block]) -> Float[Array, '']

Evaluate the energy function of the EBM given some state information.

Arguments:

• state: The state for which to evaluate the energy function. Must be compatible with blocks.
• blocks: Specifies how the information in state is organized. May be either a pre-built BlockSpec (fast path — avoids rebuilding the spec) or a plain list[Block] for convenience when calling from user code.

Returns:

A scalar representing the energy value associated with state.

with_beta(beta: Array) -> AbstractEBM

Return a copy of this EBM with a different inverse-temperature β.

Subclasses that want to work with tune_schedule(ebm=..., program=...) must override this method.

BlockSamplingProgram

A PGM block-sampling program.

This class encapsulates everything that is needed to run a PGM block sampling program in hamon. per_block_interactions and per_block_interaction_active are parallel to the free blocks in gibbs_spec, and their members are passed directly to a sampler when the state of the corresponding free block is being updated during a sampling program. per_block_interaction_global_inds and per_block_interaction_global_slices are also parallel to the free blocks, and are used to slice the global state of the program to produce the state information required to update the state of each block alongside the static information contained in the interactions.

Attributes:

• gibbs_spec: A division of some PGM into free and clamped blocks.
• samplers: A sampler to use to update every free block in gibbs_spec.
• per_block_interactions: All the interactions that touch each free block in gibbs_spec.
• per_block_interaction_active: indicates which interactions are real and which interactions are not part of the model and have been added to pad data structures so that they can be rectangular.
• per_block_interaction_global_inds: how to find the information required to update each block within the global state list
• per_block_interaction_global_slices: how to slice each array in the global state list to find the information required to update each block
• _block_sd_inds: precomputed sd_index for each free block (avoids recomputing inside scan)
• _block_positions: precomputed node positions in global state for each free block (avoids recomputing inside scan)
• _block_output_sds: precomputed output ShapeDtypeStruct pytree for each free block
• _block_slice_starts: static start index when the block occupies a contiguous range of the global state (always the case for blocks laid out by BlockSpec), or None to fall back to a gather-index scatter. A contiguous range lets the write-back lower to lax.dynamic_update_slice instead of a scatter, which XLA fuses far better.

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__init__(gibbs_spec: BlockGibbsSpec, samplers: list[hamon.conditional_samplers.AbstractConditionalSampler], interaction_groups: list[hamon.interaction.InteractionGroup])

Construct a BlockSamplingProgram.

Takes in a set of information that implicitly defines a sampling program and manipulates it into a shape appropriate for vectorized block-sampling. This involves reindexing, slicing, and often padding.

Arguments:

• gibbs_spec: A division of some PGM into free and clamped blocks.
• samplers: The update rule to use for each free block in gibbs_spec.
• interaction_groups: A list of InteractionGroups that define how the variables in your sampling program affect one another.

with_ebm(ebm) -> BlockSamplingProgram

Return a copy of this program rewired to a different EBM.

Subclasses that want to work with tune_schedule(ebm=..., program=...) must override this method.

autotune(key: jax.Array, *, ebm: AbstractEBM | None = None, program: BlockSamplingProgram | None = None, ebm_factory: Callable | None = None, program_factory: Callable | None = None, init_factory: Callable, clamp_state: list | None = None, sample_nodes: Sequence | None = None, beta_range: tuple[float, float] = (0.0, 1.0), target_acceptance: float = 0.5, min_chains: int = 3, max_chains: int = 128, initial_n: int | None = None, seed_from_energy: bool = True, gibbs_steps_per_round: int | None = None, search_exploration: bool = False, max_exploration_steps: int = 8, cost_model: bool = True, select_by: str = 'cost', target_efficiency: float = 0.9, rounds_per_probe: int = 400, n_tune: int = 4, n_polish: int = 2, n_rounds: int = 1000, compile_cache: bool | str = True, pad_probes: bool | None = None, device: DeviceLike = 'auto') -> NRPTPlan

Autotune the full NRPT configuration: N, exploration count, and schedule.

The one-call front door for solving a problem with hamon. Runs the dependency-ordered, cheap→expensive recipe and returns an :class:NRPTPlan you draw from with :meth:NRPTPlan.sample:

1. N via :func:tune_chains, probed at the final n_expl when it is already known (pinned or the deterministic device default) so stage 3 reuses stage 1's compiled round loop — the biggest cold-run compile. When the n_expl search is on, probes run at n_expl=1 instead (cheapest; Λ — hence N* — is invariant to n_expl).
2. n_expl — by default a deterministic device-calibrated count (accelerator → a fixed mid-range value, CPU → 1): reproducible across runs and ~free, since the ESS-per-wall-second objective is flat in n_expl on a dispatch-bound accelerator. Pin it explicitly with gibbs_steps_per_round (e.g. a value calibrated for your hardware), or pass search_exploration=True to tune it via :func:tune_exploration at the fixed N, reusing the schedule from step 1 (the equi-acceptance schedule is n_expl-invariant, so this needs no re-tuning and never re-discovers N); the "cost" search maximizes ESS per measured wall-second but its pick is not reproducible across runs (it depends on the machine's clock state), so it is best used as a one-time per-hardware calibration.
3. Schedule polish via :func:tune_schedule at the chosen (N, n_expl), which also leaves an equilibrated warm cold-chain state.

The multi-probe search recompiles per chain count and per n_expl, so by default the persistent compilation cache is enabled (compile_cache) to amortize those compiles across probes and runs.

Pass either a single template ebm + program (temperature-linear mode) or per-chain ebm_factory + program_factory, exactly as the individual tuners accept. init_factory(n_chains, ebms, programs) -> list builds one initial state per chain at the discovered N.

Parameters:

Name	Type	Description	Default
key	Array	PRNG key.	required
ebm / program		single template objects (temperature-linear mode), or	required
ebm_factory / program_factory		per-chain factories.	required
init_factory	Callable	(n_chains, ebms, programs) -> list of initial states.	required
clamp_state	list | None	clamped block states.	None
sample_nodes	Sequence | None	nodes defining the column order of drawn samples (must be free nodes of the program). None (default) uses all free nodes in free-block order; pass the model's canonical node list to get samples in that order (single node type only).	None
beta_range	tuple[float, float]	(β_min, β_max) temperature range.	(0.0, 1.0)
target_acceptance	float	per-pair swap acceptance target for the N search. Default 0.5 — the round-trip-optimal r = 1/2 (N ≈ 2Λ; Syed et al.).	0.5
min_chains / max_chains / initial_n		N-search bounds / start.	required
seed_from_energy	bool	seed the chain-count search from a cheap energy-variance Λ̂ (no PT ladder) so it converges in one probe — but only when local exploration mixes; a Gelman–Rubin R̂ check falls back to the robust max_chains pilot on glassy targets, so it never under-provisions. Same discovered N, fewer compiles when it applies. Default True; pass False to always run the pilot; see :func:hamon.tuning.tune_chains.	True
gibbs_steps_per_round	int | None	pin n_expl to this value, skipping both the device default and the search (step 2). For hardware you have already calibrated. None (default) uses the device default or the search.	None
search_exploration	bool	tune n_expl by a wall-timed search (step 2). Default False uses a deterministic device-calibrated n_expl (reproducible across runs); True runs :func:tune_exploration. Ignored when gibbs_steps_per_round is set.	False
max_exploration_steps	int	ceiling for the n_expl doubling search (when search_exploration=True).	8
select_by	str	for search_exploration=True — "cost" (default) maximizes cold-chain ESS per wall-second; "ele" picks n_expl by the deterministic round-trip efficiency knee (reproducible, but optimizes index-process mixing rather than sample ESS). See :func:tune_exploration.	'cost'
target_efficiency	float	ELE-adequacy threshold for select_by="ele".	0.9
cost_model	bool	for the select_by="cost" path, fit one n_expl cost line from reused production timings instead of timing each probe separately; see :func:tune_exploration.	True
rounds_per_probe	int	rounds per tuning/exploration probe (the cheap search budget).	400
n_tune	int	schedule-tuning phases per N probe.	4
n_polish	int	schedule-tuning phases in the final polish.	2
n_rounds	int	rounds for the final production run — equilibrates the warm cold state and is the window the reported round-trip rate / efficiency are measured over. Should be ≫ 2·N for a representative rate; the default (1000) suits the autotuned chain counts.	1000
compile_cache	bool | str	True enables the persistent compile cache at the default path, a str enables it at that path, False leaves placement untouched. See :func:hamon.enable_persistent_compile_cache.	True
pad_probes	bool | None	chain-mask the stage-1 probes — pad every probe's round loop to max_chains so all probes share ONE compiled loop instead of recompiling per chain count (see :func:hamon.tuning.tune_chains). None (default) enables it on an accelerator in template mode (where the padding Gibbs work is ~free) and disables it on CPU or the factory route. Stages 2-3 and the production draw always run unpadded.	None
device	DeviceLike	where to run; resolved once and reused across every stage.	'auto'

Returns:

Name	Type	Description
An	NRPTPlan	class:NRPTPlan.

autosample(key: jax.Array, *, n_samples: int, n_warmup: int = 0, steps_per_sample: int = 1, **autotune_kwargs) -> tuple[jax.Array, AutotuneReport]

One-shot: :func:autotune then draw — returns (samples, report).

The convenience entry for "give me samples." Forwards all keyword arguments to :func:autotune (ebm/program or factories, init_factory, beta_range, device, …), then draws n_samples from the tuned plan. For repeated draws from one tuned configuration, call :func:autotune once and reuse :meth:NRPTPlan.sample.

resolve_entry_device(device: DeviceLike, *, n_chains: int, n_nodes: int, arrays: Any = ()) -> JaxDevice | None

Entry-point resolution: heuristic score plus a tracer guard.

When any entry array is a tracer the caller is already inside a jit/vmap/grad trace, where opening device contexts or transferring arrays is not meaningful — routing becomes a no-op.

tune_chains(key: jax.Array, ebm_factory: Callable | None = None, program_factory: Callable | None = None, init_factory: Callable | None = None, clamp_state: list | None = None, beta_range: tuple[float, float] = (0.0, 1.0), gibbs_steps_per_round: int = 0, initial_n: int | None = None, seed_from_energy: bool = True, target_acceptance: float = 0.5, rounds_per_probe: int = 200, n_tune_per_probe: int = 4, max_iters: int = 6, min_chains: int = 3, max_chains: int = 128, lambda_rtol: float = 0.05, *, ebm: AbstractEBM | None = None, program: BlockSamplingProgram | None = None, tune_tol: float | None = None, safety_margin: float = 0.05, device: DeviceLike = 'auto', pad_probes: bool = False) -> dict

Iteratively discover the right chain count for a given target acceptance.

Follows the N-tuning method of Syed et al. (2021): the global communication barrier Λ is a schedule invariant (Σ rejection_rates ≈ Λ at any chain count), so it is estimated at a single fixed N from a schedule-tuned run rather than searched for by probing many chain counts.

1. Estimate Λ̂ = Σ rejection_rates at the current N (each probe runs tune_schedule, which tunes the schedule toward equi-acceptance).
2. Recommend N = ceil(Λ̂·(1 + safety_margin) / r_target) + 1 — the round-trip-optimal 2Λ + 1 chains at r = 1/2 (target_acceptance = 0.5).
3. Iterate this fixed point (re-estimate Λ̂ at N) until N stops moving.

Because Λ̂ comes from the current probe (not a running maximum), the result is essentially independent of the starting N — discovery from initial_n=None and from a reasonable guess converge to the same count. With no initial_n the first probe runs at a high pilot of max_chains chains: high on purpose, because a low pilot's rejection rates saturate and bias Λ̂ low, forcing the fixed point to climb over several probes. An over-resolved pilot gives an unbiased Λ̂ in one probe, landing n* within ±1 immediately, so discovery converges in ~2 probes regardless of problem size.

Instead of providing ebm_factory and program_factory, you can pass a template ebm and program and the factories will be built internally using ebm.with_beta() and program.with_ebm(). init_factory is still required as initialization varies by use case.

Parameters:

Name	Type	Description	Default
key	Array	PRNG key	required
ebm_factory	Callable | None	betas_array → list[EBM]	None
program_factory	Callable | None	list[EBM] → list[Program]	None
init_factory	Callable | None	(n_chains, list[EBM], list[Program]) → list[init_states]. Receives EBMs and programs so it can extract the correct free_blocks for initialization (block nodes must be the same objects as the EBMs' nodes).	None
clamp_state	list | None	clamped block states	None
beta_range	tuple[float, float]	(β_min, β_max) for the temperature range	(0.0, 1.0)
gibbs_steps_per_round	int	Gibbs sweeps between swap attempts	0
initial_n	int | None	starting chain count. The default None runs a high pilot probe at max_chains for an unbiased Λ̂ (no initial guess needed), unless seed_from_energy is set; pass an int to start there instead.	None
seed_from_energy	bool	seed the search from a cheap energy-variance Λ̂ (Theorem 2, no PT ladder; see :func:_estimate_barrier_energy) so it converges in one probe instead of the max_chains pilot's two — fewer compiles. Default True (applies when initial_n is None); pass False to always run the pilot. Self-guarding: the estimate is only trustworthy when local exploration mixes, so the energy probe also returns a Gelman–Rubin R̂; if R̂ exceeds the cutoff (trapping — a glassy target where the estimate would be unreliable) the search falls back to the robust max_chains pilot. On a mixing target (R̂≈1) the seed lands on N and, because the probe RNG is key-aligned with the pilot, the discovered N and schedule are bit-identical to the pilot path; on a glassy target it is exactly the pilot. So it never under-provisions — it only ever saves (mixing) or matches* the pilot (glassy), at the cost of the energy probe.	True
target_acceptance	float	desired per-pair swap acceptance rate. Default 0.5 — the round-trip-optimal rejection r = 1/2 (N ≈ 2Λ; Syed et al.), not the 0.77 from the reversible-PT literature.	0.5
rounds_per_probe	int	rounds per probe (and for the final production probe)	200
n_tune_per_probe	int	schedule tuning iterations for the final probe	4
max_iters	int	maximum discovery iterations	6
min_chains	int	floor on chain count	3
max_chains	int	ceiling on chain count	128
lambda_rtol	float	relative tolerance for Λ stabilization (default 5%)	0.05
safety_margin	float	small fractional pad on N* (default 0.05) covering residual barrier bias and ELE-assumption violations; 0.0 gives the bare round-trip-optimal count	0.05
pad_probes	bool	run every probe's round loop padded to max_chains with chain masking (see :func:hamon.nrpt.nrpt), so probes at different chain counts share ONE compiled round loop instead of recompiling per count — the dominant cold cost of discovery. Padding chains do wasted-but-decoupled Gibbs work (~free on a dispatch-bound accelerator; real cost on CPU, so leave off there). Template (temperature-linear) mode only. Probe statistics are computed on the sliced live prefix, but the probe RNG stream differs from an unpadded run, so discovered N can shift within its normal probe-to-probe variability.	False

Returns:

Type

Description

dict

dict with keys: n_chains: final recommended chain count betas: optimized schedule at that chain count Lambda: conservative (max) barrier estimate Lambda_raw: last raw estimate (may be lower than Lambda) target_acceptance: the target used converged_reason: "chain_count" | "lambda_stable" | "max_iters" barrier_identified: whether the final count's ladder round-tripped, so the reported Lambda is identified rather than a stalled-conveyor within-basin artifact (None if round trips were not tracked). history: list of per-probe dicts (each carries barrier_identified)

tune_exploration(key: jax.Array, ebm_factory: Callable | None = None, program_factory: Callable | None = None, init_states: Sequence[list] = (), clamp_state: list | None = None, initial_betas: jax.Array | None = None, start_steps: int = 1, max_steps: int = 64, rounds_per_probe: int = 400, n_tune_per_probe: int = 3, improve_tol: float = 0.05, time_rounds: int = 200, time_reps: int = 3, cost_model: bool = True, select_by: str = 'cost', target_efficiency: float = 0.9, *, fixed_schedule: jax.Array | None = None, ebm: AbstractEBM | None = None, program: BlockSamplingProgram | None = None, device: DeviceLike = 'auto') -> dict

Discover the local-exploration count gibbs_steps_per_round (n_expl).

n_expl is the only major NRPT knob hamon does not otherwise auto-tune (tune_chains sets N from Λ; tune_schedule sets the schedule). The objective maximized here is effective sample size per measured steady-state wall-second,

objective(n_expl) = ESS_median(cold chain) / (rounds · t_round(n_expl)),

where t_round is the per-round wall time measured on the target device after warm-up (so XLA compile is excluded). This is the honest endgame and, crucially, it self-calibrates to the hardware's real cost structure. The per-round cost is t_round = c₀ + n_expl·c_s: a fixed overhead c₀ (swap pass, energy recompute, host/kernel dispatch, scan bookkeeping) plus n_expl Gibbs sweeps at c_s each. The Syed et al. compute model assumes c₀ = 0 (cost ∝ n_expl), which makes n_expl=1 optimal; but cold-chain ESS grows sub-linearly in n_expl while real per-round cost grows less than linearly when c₀ is non-trivial — so on a dispatch-bound backend where c₀ ≳ 1.4·c_s the optimum shifts to n_expl > 1. Measuring t_round rather than assuming cost ∝ n_expl is what lets the search see this.

Why not the cheaper round-trip proxy τ_obs / n_expl? Empirically it under-picks n_expl: round trips count excursions but not how decorrelated each cold sample is, so doubling n_expl can sharply raise cold-chain ESS while barely moving τ_obs (the r≈0.81 ESS↔round-trip correlation of Syed et al. breaking down). The round-trip rate is still used as a robust gate and cross-check — efficiency = τ_obs/τ̄ is the ELE-violation meter, and the efficiency_limiter from :func:report_nrpt_diagnostics stops the search when a probe is schedule-limited (an unequalized ladder, where more local exploration cannot help). rt_per_compute and t_round are recorded per probe alongside the ESS-per-second objective.

See :func:tune_chains for the N analogue. The chain count is held fixed at len(initial_betas) — Λ (hence N) is a schedule invariant robust to n_expl, so the two searches decouple; run tune_chains first if N is unknown.

Each probe (1) gets a schedule, (2) measures ESS over a cold-chain trace via an :class:~hamon.NRPTStateObserver, and (3) times the steady-state round loop (warm-up absorbs the one-time compile; the median of time_reps runs of time_rounds rounds gives t_round). With fixed_schedule=None (default) the schedule is re-tuned per probe via :func:tune_schedule; pass fixed_schedule (a pre-tuned ladder) to reuse it and run each probe as a single :func:nrpt production call — much cheaper, and sound because the equi-acceptance schedule is invariant to n_expl. autotune uses this mode. Instead of ebm_factory / program_factory you may pass a single template ebm and program (temperature-linear mode).

Parameters:

Name	Type	Description	Default
key	Array	PRNG key.	required
ebm_factory / program_factory		per-chain factories, or use ebm / program template objects.	required
init_states	Sequence[list]	one initial block-state list per chain (fixed across probes).	()
clamp_state	list | None	clamped block states.	None
initial_betas	Array | None	the (fixed) β ladder; its length sets the chain count.	None
start_steps	int	smallest n_expl to try (≥ 1).	1
max_steps	int	largest n_expl before the search stops.	64
rounds_per_probe	int	production rounds per probe (and the tuning ceiling) — must be large enough for ESS and τ_obs to be low-variance.	400
n_tune_per_probe	int	schedule-tuning phases per probe.	3
improve_tol	float	minimum fractional objective gain to keep doubling (guards against chasing Monte-Carlo noise past the peak).	0.05
time_rounds	int	rounds per timed run when measuring t_round.	200
time_reps	int	number of timed runs to reduce over (noise control).	3
select_by	str	how to choose n_expl. "cost" (default) maximizes cold-chain ESS per wall-second (see cost_model) — the sample-quality objective; the pick depends on the machine's measured cost ratio, so it is best used as a one-time per-hardware calibration. "ele" instead picks the smallest count whose round-trip efficiency τ_obs/τ̂ reaches the ELE-adequacy knee — deterministic (no wall-clock) and the criterion the Syed et al. analysis prescribes, but it optimizes index-process mixing rather than cold-sample ESS, so it under-picks n_expl on a dispatch-bound accelerator where extra sweeps are nearly free; use it for index-efficiency or severe-ELE-violation regimes.	'cost'
target_efficiency	float	ELE-adequacy threshold for select_by="ele" — stop climbing once τ_obs/τ̂ reaches this (it also stops on a plateau or a schedule-limited verdict).	0.9
cost_model	bool	for select_by="cost" only — when True (and a fixed_schedule is given) fit t_round = c₀ + n_expl·c_s by least squares across the probes (each timing reuses the production executable, no separate observer=None compile) and take the argmax from that shared line; False times each probe independently (the flat objective then lets timing noise pick the count at random).	True
fixed_schedule	Array | None	a pre-tuned β ladder to reuse across all probes (each probe becomes one production run, no per-probe re-tuning). None (default) re-tunes per probe; autotune passes the ladder from tune_chains.	None
device	DeviceLike	where to run; resolved once and reused across probes. Timing is measured on this device, so the chosen n_expl is calibrated to it.	'auto'

Returns:

Type

Description

dict

dict with keys: gibbs_steps_per_round: the chosen n_expl. objective: ESS per measured wall-second at the chosen n_expl. ess_median / tau_observed / efficiency / rt_per_compute / t_round: at the choice (t_round in seconds). betas: the tuned schedule at the chosen n_expl. history: list of per-probe records (n_expl, objective, ess_median, tau_obs, rt_per_compute, t_round, efficiency, efficiency_limiter, betas).

tune_schedule(key: jax.Array, ebm_factory: Callable | None = None, program_factory: Callable | None = None, init_states: Sequence[list] = (), clamp_state: list | None = None, n_rounds: int = 0, gibbs_steps_per_round: int = 0, initial_betas: jax.Array | None = None, n_tune: int = 5, rounds_per_tune: int = 200, track_round_trips: bool = True, *, ebm: AbstractEBM | None = None, program: BlockSamplingProgram | None = None, observer: AbstractNRPTObserver | None = None, adaptive_tuning: bool = True, tune_tol: float | None = None, equalize_tol: float = 0.05, phase_patience: int = 2, min_tune_phases: int = 1, round_batch: int = 50, min_rounds_per_tune: int = 50, round_stable_k: int = 2, lambda_rtol: float = 0.05, device: DeviceLike = 'auto', pad_chains_to: int | None = None) -> tuple[list, dict]

NRPT with iterative schedule optimization (Algorithm 4).

pad_chains_to enables chain masking in every phase's round loop (see :func:hamon.nrpt.nrpt): phases at different chain counts padded to the same length share one compiled executable. Stats and states are sliced back to the true count before any tuning math sees them, so the schedule optimization is untouched. Incompatible with observer.

Adapts the β schedule over tuning phases, then runs the final n_rounds production phase with the optimized schedule. Each phase logs one INFO line (Λ, mean acceptance, schedule movement) so long runs are not silent.

Instead of providing ebm_factory and program_factory, you can pass a template ebm and program and the factories will be built internally using ebm.with_beta() and program.with_ebm().

Convergence-driven tuning (default, adaptive_tuning=True): budgets are chosen automatically, so callers need not guess them. Each phase runs as many rounds as needed for the Λ estimate to settle (between min_rounds_per_tune and the rounds_per_tune ceiling, in round_batch increments — see _tune_phase_adaptive_rounds), giving a low-variance rejection-rate estimate. The schedule with the best (lowest-spread) rejection rates seen across phases is kept for production — not the last, which can be noisier. Tuning stops once the schedule is well-equalized (std(rejection_rates) < equalize_tol) OR has settled (max|Δβ| below the effective tune_tol — its Monte-Carlo floor) for phase_patience consecutive phases, after at least min_tune_phases, capped at n_tune. (max|Δβ| alone is not a reliable convergence signal: it plateaus at a problem-dependent noise floor rather than going to zero, so the equalization check is what stops already-good schedules.) When tune_tol is left None it defaults to _DEFAULT_TUNE_TOL here. Counts are deterministic for a given seed but problem-dependent — do not assume a fixed round/phase count.

Legacy mode (adaptive_tuning=False): runs exactly n_tune phases of exactly rounds_per_tune rounds and uses the last schedule. tune_tol then behaves as the optional early-stop it always was (None ⇒ run all n_tune phases). n_tune and rounds_per_tune act as safety caps in both modes.

Returns (states, stats) where stats includes tuning history in stats["tuning_history"] (each entry records max_beta_shift). States are ordered by ascending β — the cold chain (target distribution) is states[-1].

device is resolved once here and passed to every tuning and production phase, so the device never flips mid-run (see hamon.device).

hamon.NRPTPlan

A tuned NRPT configuration plus a warm cold-chain state.

Returned by :func:autotune. Holds the discovered hyperparameters (N, schedule, n_expl) and an equilibrated cold-chain state, so :meth:sample can draw repeatedly and cheaply — no re-tuning, reusing the compiled loop.

Attributes:

Name	Type	Description
n_chains	/ betas / gibbs_steps_per_round / Lambda	the tuned config.
device	/ betas / gibbs_steps_per_round / Lambda	the resolved device (or None).
report	/ betas / gibbs_steps_per_round / Lambda	the :class:AutotuneReport.

sample(key: jax.Array, n_samples: int, *, n_warmup: int = 0, steps_per_sample: int = 1) -> jax.Array

Draw n_samples from the target (cold chain) — cheap and repeatable.

Runs single-chain block Gibbs at the tuned cold β from the stored warm state (the established post-NRPT draw). Returns a (n_samples, n_nodes) array; call again with a fresh key for more, with no re-tuning.

hamon.AutotuneReport

Diagnostics from an :func:autotune run.

Attributes:

Name	Type	Description
n_chains		discovered chain count N.
gibbs_steps_per_round		discovered local-exploration count n_expl.
Lambda		estimated global communication barrier Λ.
betas		the final tuned β ladder.
device		the resolved device (string) or None.
chain_history		per-probe records from the N search (:func:tune_chains).
exploration		the :func:tune_exploration result dict, or None when the n_expl search was skipped.
round_trip_diagnostics		round-trip summary from the final production run.
total_round_trips		total completed round trips observed during the final production run (summed across chains), or None.
production_rounds		number of rounds the production run used (the window total_round_trips and tau_observed were measured over).

summary() -> str

Human-readable multi-line summary.

Building blocks

hamon.nrpt

Non-Reversible Parallel Tempering with vectorized swaps.

Based on Syed et al. (2021), "Non-Reversible Parallel Tempering: a Scalable Highly Parallel MCMC Scheme" (arXiv:1905.02939).

Exploits temperature-linearity (E_β = β·E_base) for single-eval-per-chain swap decisions. Adaptive schedule optimization (Algorithm 4) equalizes rejection rates. Optional energy caching with boundary-only deltas for rectangular block partitions.

DeviceLike module-attribute

Represent a PEP 604 union type

E.g. for int | str

__cached__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__doc__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__file__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__name__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__package__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

annotations module-attribute

logger module-attribute

Instances of the Logger class represent a single logging channel. A "logging channel" indicates an area of an application. Exactly how an "area" is defined is up to the application developer. Since an application can have any number of areas, logging channels are identified by a unique string. Application areas can be nested (e.g. an area of "input processing" might include sub-areas "read CSV files", "read XLS files" and "read Gnumeric files"). To cater for this natural nesting, channel names are organized into a namespace hierarchy where levels are separated by periods, much like the Java or Python package namespace. So in the instance given above, channel names might be "input" for the upper level, and "input.csv", "input.xls" and "input.gnu" for the sub-levels. There is no arbitrary limit to the depth of nesting.

NRPTCarry

Scan carry for the NRPT inner loop.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__orig_bases__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__slots__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

accepted class-attribute

Alias for field number 2

attempted class-attribute

Alias for field number 3

base_E class-attribute

Alias for field number 5

idx_state class-attribute

Alias for field number 4

key class-attribute

Alias for field number 0

obs_carry class-attribute

Alias for field number 6

states class-attribute

Alias for field number 1

AbstractEBM

Something that has a well-defined energy function (map from a state to a scalar).

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__init__() -> None

Initialize self. See help(type(self)) for accurate signature.

energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], _State]], blocks: BlockSpec | list[Block]) -> Float[Array, '']

Evaluate the energy function of the EBM given some state information.

Arguments:

• state: The state for which to evaluate the energy function. Must be compatible with blocks.
• blocks: Specifies how the information in state is organized. May be either a pre-built BlockSpec (fast path — avoids rebuilding the spec) or a plain list[Block] for convenience when calling from user code.

Returns:

A scalar representing the energy value associated with state.

with_beta(beta: Array) -> AbstractEBM

Return a copy of this EBM with a different inverse-temperature β.

Subclasses that want to work with tune_schedule(ebm=..., program=...) must override this method.

AbstractNRPTObserver

Observer for NRPT rounds, called once per round after Gibbs sweeps and swaps.

Concrete subclasses must implement __call__ and may override init to provide a non-trivial carry. The observation returned by __call__ is stacked by lax.scan into a pytree with a leading axis of size n_rounds. Return None as the observation for accumulate-only observers that do not need per-round storage.

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__call__(stacked_states: list[Array], base_energies: Array, round_idx: Int[Array, ''], carry: ~ObserveCarry) -> tuple[~ObserveCarry, PyTree]

Observe one NRPT round.

Arguments:

• stacked_states: Per-block arrays, each of shape (n_chains, ...). The cold chain (target) is at index -1.
• base_energies: Shape (n_chains,) base energies (no β factor), aligned with stacked_states — base_energies[c] is the energy of the state at chain position c after this round's swaps.
• round_idx: Zero-based round counter.
• carry: Arbitrary pytree state threaded across rounds.

Returns:

(updated_carry, observation) — observation is stacked by lax.scan; use None for accumulate-only mode.

__init__() -> None

Initialize self. See help(type(self)) for accurate signature.

init() -> PyTree

Initialize the observer carry. Defaults to None.

BlockSamplingProgram

A PGM block-sampling program.

This class encapsulates everything that is needed to run a PGM block sampling program in hamon. per_block_interactions and per_block_interaction_active are parallel to the free blocks in gibbs_spec, and their members are passed directly to a sampler when the state of the corresponding free block is being updated during a sampling program. per_block_interaction_global_inds and per_block_interaction_global_slices are also parallel to the free blocks, and are used to slice the global state of the program to produce the state information required to update the state of each block alongside the static information contained in the interactions.

Attributes:

• gibbs_spec: A division of some PGM into free and clamped blocks.
• samplers: A sampler to use to update every free block in gibbs_spec.
• per_block_interactions: All the interactions that touch each free block in gibbs_spec.
• per_block_interaction_active: indicates which interactions are real and which interactions are not part of the model and have been added to pad data structures so that they can be rectangular.
• per_block_interaction_global_inds: how to find the information required to update each block within the global state list
• per_block_interaction_global_slices: how to slice each array in the global state list to find the information required to update each block
• _block_sd_inds: precomputed sd_index for each free block (avoids recomputing inside scan)
• _block_positions: precomputed node positions in global state for each free block (avoids recomputing inside scan)
• _block_output_sds: precomputed output ShapeDtypeStruct pytree for each free block
• _block_slice_starts: static start index when the block occupies a contiguous range of the global state (always the case for blocks laid out by BlockSpec), or None to fall back to a gather-index scatter. A contiguous range lets the write-back lower to lax.dynamic_update_slice instead of a scatter, which XLA fuses far better.

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__init__(gibbs_spec: BlockGibbsSpec, samplers: list[hamon.conditional_samplers.AbstractConditionalSampler], interaction_groups: list[hamon.interaction.InteractionGroup])

Construct a BlockSamplingProgram.

Takes in a set of information that implicitly defines a sampling program and manipulates it into a shape appropriate for vectorized block-sampling. This involves reindexing, slicing, and often padding.

Arguments:

• gibbs_spec: A division of some PGM into free and clamped blocks.
• samplers: The update rule to use for each free block in gibbs_spec.
• interaction_groups: A list of InteractionGroups that define how the variables in your sampling program affect one another.

with_ebm(ebm) -> BlockSamplingProgram

Return a copy of this program rewired to a different EBM.

Subclasses that want to work with tune_schedule(ebm=..., program=...) must override this method.

optimize_schedule(rejection_rates: jax.Array, betas: jax.Array) -> jax.Array

Equalize per-pair rejection rates by redistributing β values.

Estimates the cumulative communication barrier Λ(β_k) = Σ_{i≤k} r_i, then places the new ladder at a regular grid of Λ (the equi-acceptance schedule of Syed et al. 2021, Algorithm 2: β_k = Λ⁻¹(k/N · Λ)). The inverse Λ⁻¹ is taken as a Fritsch–Carlson monotone cubic* of β against the cumulative barrier — smoother than the previous piecewise-linear inverse while staying monotone (no overshoot), which the paper recommends over linear interpolation.

The result keeps the dtype of betas so repeated tuning phases do not drift to float64 when x64 is enabled.

nrpt(key: jax.Array, ebms: Sequence[AbstractEBM] | AbstractEBM, programs: Sequence[BlockSamplingProgram] | BlockSamplingProgram, init_states: Sequence[list], clamp_state: list, n_rounds: int, gibbs_steps_per_round: int, betas: jax.Array | None = None, track_round_trips: bool = True, energy_delta_fn: Callable | None = None, observer: AbstractNRPTObserver | None = None, device: DeviceLike = 'auto', pad_chains_to: int | None = None, _emit_diagnostics: bool = True, _return_stacked: bool = False) -> tuple[list, dict]

Non-Reversible Parallel Tempering with vectorized swaps.

pad_chains_to (≥ n_chains) enables chain masking: the ladder is padded to that fixed length with copies of the coldest chain and the round loop runs with the true chain count as traced data — swaps beyond the live prefix are forced-rejected (identity permutation), so the padding is fully decoupled from the live ladder. All returned states/stats are sliced back to the true count, so callers see exactly the shapes and semantics of an unpadded run. The point: probes at different chain counts padded to the same length share ONE compiled round loop instead of recompiling per count (the dominant cold cost of tune_chains), at the price of wasted Gibbs work on the padding chains (~free on a dispatch-bound accelerator). Temperature-linear mode only; incompatible with observer and energy_delta_fn.

Single-pass DEO: one swap parity per round, alternating even/odd. Multi-pass breaks non-reversibility (even∘odd∘odd∘even = identity).

Temperature-linear mode: instead of per-chain sequences, ebms and programs may each be a single template object (any β; it is rebased to β = 1 via with_beta()/with_ebm()). betas is then required and defines the chain count. Interactions are scaled by each chain's β inside the sampling kernel, which assumes every interaction array is linear in β — true for the DiscreteEBMFactor family (and anything built from β-scaled factor weights), and consistent with the E_β = β·E_base assumption the swap math already makes. This avoids constructing one program per chain and storing per-chain copies of every interaction tensor. For models whose interactions are not linear in β, pass explicit per-chain sequences instead.

Chains are ordered by ascending β: index 0 is the hottest chain (lowest β, closest to the reference distribution) and index −1 is the coldest chain (highest β, the target distribution you want to sample from). The returned states list preserves this ordering, and betas must be sorted ascending (validated).

init_states may be either a sequence of per-chain block-state lists, or a single block-state list whose arrays carry a leading (n_chains, ...) axis — e.g. straight from hinton_init(key, model, blocks, (n_chains,)) — avoiding the per-chain list/restack dance. A hottest chain at exactly β = 0 (sampling the reference distribution) is supported; base energies are computed from a β = 1 copy of the EBM when with_beta() is available, falling back to the coldest chain.

.. warning:: To collect samples from the target distribution, always use states[-1] (the cold chain), not states[0].

device selects where the computation runs: "auto" (default) routes small workloads to the CPU and large ones to a visible accelerator (see hamon.device), "cpu"/"gpu" force a platform, a concrete jax.Device is used as-is, and None leaves placement untouched. Routing re-commits the entry arrays (program tensors, states, betas) to the chosen device; outputs come back committed there. Arrays closed over by energy_delta_fn cannot be moved by routing — keep them uncommitted or on the target device.

Stats keys

accepted, attempted, acceptance_rate, rejection_rates, betas round_trip_diagnostics (if track_round_trips=True): Lambda, tau_predicted, tau_observed, efficiency, lambda_profile, round_trips_per_chain, restarts_per_chain observations (if observer is not None): Per-round observer output stacked along axis 0. observer_carry (if observer is not None): Final observer carry after all rounds.

free_node_count(program) -> int

Total free nodes in a sampling program — O(1) Python metadata.

init_index_state(n_chains: int) -> dict

Initialize index process tracking arrays.

machine_to_chain[j] = which chain position machine j's state currently occupies. Initially machine j is at chain j.

visited_top[j] = whether machine j has reached chain N since its last round trip completion.

Returns a dict suitable for inclusion in lax.scan carry.

resolve_entry_device(device: DeviceLike, *, n_chains: int, n_nodes: int, arrays: Any = ()) -> JaxDevice | None

Entry-point resolution: heuristic score plus a tracer guard.

When any entry array is a tracer the caller is already inside a jit/vmap/grad trace, where opening device contexts or transferring arrays is not meaningful — routing becomes a no-op.

round_trip_summary(index_state: dict, rejection_rates: jax.Array, betas: jax.Array, n_rounds: int) -> dict

Compute full diagnostic summary for NRPT run.

Jitted so the handful of reductions below (Λ, τ̄, the local-barrier profile, the round-trip rate) fuse into a single compiled kernel instead of ~8 eager op-by-op dispatches, each of which otherwise pays a first-shape XLA compile when called once per probe at a new chain count. n_rounds is traced (not static), so the compile is shared across round counts.

Returns dict with

Lambda: global communication barrier estimate tau_predicted: theoretical optimal round trip rate tau_observed: empirical round trip rate efficiency: tau_observed / tau_predicted (closer to 1 = better) lambda_profile: local barrier at each pair midpoint round_trips_per_chain: per-machine round trip counts restarts_per_chain: per-machine restart counts

tree_device_put(tree: _T, device: JaxDevice | None) -> _T

Commit every jax.Array leaf of tree to device.

Non-array leaves (blocks, nodes, samplers, specs) pass through with object identity preserved, so equinox's static partition hashes identically and jit caches stay warm. If every array leaf is already committed to device, the original object is returned unchanged — this keeps repeated calls (e.g. tune_schedule tuning phases) presenting the literally-same pytree to the jit cache.

hamon.tune_schedule(key: jax.Array, ebm_factory: Callable | None = None, program_factory: Callable | None = None, init_states: Sequence[list] = (), clamp_state: list | None = None, n_rounds: int = 0, gibbs_steps_per_round: int = 0, initial_betas: jax.Array | None = None, n_tune: int = 5, rounds_per_tune: int = 200, track_round_trips: bool = True, *, ebm: AbstractEBM | None = None, program: BlockSamplingProgram | None = None, observer: AbstractNRPTObserver | None = None, adaptive_tuning: bool = True, tune_tol: float | None = None, equalize_tol: float = 0.05, phase_patience: int = 2, min_tune_phases: int = 1, round_batch: int = 50, min_rounds_per_tune: int = 50, round_stable_k: int = 2, lambda_rtol: float = 0.05, device: DeviceLike = 'auto', pad_chains_to: int | None = None) -> tuple[list, dict]

NRPT with iterative schedule optimization (Algorithm 4).

pad_chains_to enables chain masking in every phase's round loop (see :func:hamon.nrpt.nrpt): phases at different chain counts padded to the same length share one compiled executable. Stats and states are sliced back to the true count before any tuning math sees them, so the schedule optimization is untouched. Incompatible with observer.

Adapts the β schedule over tuning phases, then runs the final n_rounds production phase with the optimized schedule. Each phase logs one INFO line (Λ, mean acceptance, schedule movement) so long runs are not silent.

Instead of providing ebm_factory and program_factory, you can pass a template ebm and program and the factories will be built internally using ebm.with_beta() and program.with_ebm().

Convergence-driven tuning (default, adaptive_tuning=True): budgets are chosen automatically, so callers need not guess them. Each phase runs as many rounds as needed for the Λ estimate to settle (between min_rounds_per_tune and the rounds_per_tune ceiling, in round_batch increments — see _tune_phase_adaptive_rounds), giving a low-variance rejection-rate estimate. The schedule with the best (lowest-spread) rejection rates seen across phases is kept for production — not the last, which can be noisier. Tuning stops once the schedule is well-equalized (std(rejection_rates) < equalize_tol) OR has settled (max|Δβ| below the effective tune_tol — its Monte-Carlo floor) for phase_patience consecutive phases, after at least min_tune_phases, capped at n_tune. (max|Δβ| alone is not a reliable convergence signal: it plateaus at a problem-dependent noise floor rather than going to zero, so the equalization check is what stops already-good schedules.) When tune_tol is left None it defaults to _DEFAULT_TUNE_TOL here. Counts are deterministic for a given seed but problem-dependent — do not assume a fixed round/phase count.

Legacy mode (adaptive_tuning=False): runs exactly n_tune phases of exactly rounds_per_tune rounds and uses the last schedule. tune_tol then behaves as the optional early-stop it always was (None ⇒ run all n_tune phases). n_tune and rounds_per_tune act as safety caps in both modes.

Returns (states, stats) where stats includes tuning history in stats["tuning_history"] (each entry records max_beta_shift). States are ordered by ascending β — the cold chain (target distribution) is states[-1].

device is resolved once here and passed to every tuning and production phase, so the device never flips mid-run (see hamon.device).

hamon.optimize_schedule(rejection_rates: jax.Array, betas: jax.Array) -> jax.Array

Equalize per-pair rejection rates by redistributing β values.

Estimates the cumulative communication barrier Λ(β_k) = Σ_{i≤k} r_i, then places the new ladder at a regular grid of Λ (the equi-acceptance schedule of Syed et al. 2021, Algorithm 2: β_k = Λ⁻¹(k/N · Λ)). The inverse Λ⁻¹ is taken as a Fritsch–Carlson monotone cubic* of β against the cumulative barrier — smoother than the previous piecewise-linear inverse while staying monotone (no overshoot), which the paper recommends over linear interpolation.

The result keeps the dtype of betas so repeated tuning phases do not drift to float64 when x64 is enabled.

hamon.tune_chains(key: jax.Array, ebm_factory: Callable | None = None, program_factory: Callable | None = None, init_factory: Callable | None = None, clamp_state: list | None = None, beta_range: tuple[float, float] = (0.0, 1.0), gibbs_steps_per_round: int = 0, initial_n: int | None = None, seed_from_energy: bool = True, target_acceptance: float = 0.5, rounds_per_probe: int = 200, n_tune_per_probe: int = 4, max_iters: int = 6, min_chains: int = 3, max_chains: int = 128, lambda_rtol: float = 0.05, *, ebm: AbstractEBM | None = None, program: BlockSamplingProgram | None = None, tune_tol: float | None = None, safety_margin: float = 0.05, device: DeviceLike = 'auto', pad_probes: bool = False) -> dict

Iteratively discover the right chain count for a given target acceptance.

Follows the N-tuning method of Syed et al. (2021): the global communication barrier Λ is a schedule invariant (Σ rejection_rates ≈ Λ at any chain count), so it is estimated at a single fixed N from a schedule-tuned run rather than searched for by probing many chain counts.

1. Estimate Λ̂ = Σ rejection_rates at the current N (each probe runs tune_schedule, which tunes the schedule toward equi-acceptance).
2. Recommend N = ceil(Λ̂·(1 + safety_margin) / r_target) + 1 — the round-trip-optimal 2Λ + 1 chains at r = 1/2 (target_acceptance = 0.5).
3. Iterate this fixed point (re-estimate Λ̂ at N) until N stops moving.

Because Λ̂ comes from the current probe (not a running maximum), the result is essentially independent of the starting N — discovery from initial_n=None and from a reasonable guess converge to the same count. With no initial_n the first probe runs at a high pilot of max_chains chains: high on purpose, because a low pilot's rejection rates saturate and bias Λ̂ low, forcing the fixed point to climb over several probes. An over-resolved pilot gives an unbiased Λ̂ in one probe, landing n* within ±1 immediately, so discovery converges in ~2 probes regardless of problem size.

Instead of providing ebm_factory and program_factory, you can pass a template ebm and program and the factories will be built internally using ebm.with_beta() and program.with_ebm(). init_factory is still required as initialization varies by use case.

Parameters:

Name	Type	Description	Default
key	Array	PRNG key	required
ebm_factory	Callable | None	betas_array → list[EBM]	None
program_factory	Callable | None	list[EBM] → list[Program]	None
init_factory	Callable | None	(n_chains, list[EBM], list[Program]) → list[init_states]. Receives EBMs and programs so it can extract the correct free_blocks for initialization (block nodes must be the same objects as the EBMs' nodes).	None
clamp_state	list | None	clamped block states	None
beta_range	tuple[float, float]	(β_min, β_max) for the temperature range	(0.0, 1.0)
gibbs_steps_per_round	int	Gibbs sweeps between swap attempts	0
initial_n	int | None	starting chain count. The default None runs a high pilot probe at max_chains for an unbiased Λ̂ (no initial guess needed), unless seed_from_energy is set; pass an int to start there instead.	None
seed_from_energy	bool	seed the search from a cheap energy-variance Λ̂ (Theorem 2, no PT ladder; see :func:_estimate_barrier_energy) so it converges in one probe instead of the max_chains pilot's two — fewer compiles. Default True (applies when initial_n is None); pass False to always run the pilot. Self-guarding: the estimate is only trustworthy when local exploration mixes, so the energy probe also returns a Gelman–Rubin R̂; if R̂ exceeds the cutoff (trapping — a glassy target where the estimate would be unreliable) the search falls back to the robust max_chains pilot. On a mixing target (R̂≈1) the seed lands on N and, because the probe RNG is key-aligned with the pilot, the discovered N and schedule are bit-identical to the pilot path; on a glassy target it is exactly the pilot. So it never under-provisions — it only ever saves (mixing) or matches* the pilot (glassy), at the cost of the energy probe.	True
target_acceptance	float	desired per-pair swap acceptance rate. Default 0.5 — the round-trip-optimal rejection r = 1/2 (N ≈ 2Λ; Syed et al.), not the 0.77 from the reversible-PT literature.	0.5
rounds_per_probe	int	rounds per probe (and for the final production probe)	200
n_tune_per_probe	int	schedule tuning iterations for the final probe	4
max_iters	int	maximum discovery iterations	6
min_chains	int	floor on chain count	3
max_chains	int	ceiling on chain count	128
lambda_rtol	float	relative tolerance for Λ stabilization (default 5%)	0.05
safety_margin	float	small fractional pad on N* (default 0.05) covering residual barrier bias and ELE-assumption violations; 0.0 gives the bare round-trip-optimal count	0.05
pad_probes	bool	run every probe's round loop padded to max_chains with chain masking (see :func:hamon.nrpt.nrpt), so probes at different chain counts share ONE compiled round loop instead of recompiling per count — the dominant cold cost of discovery. Padding chains do wasted-but-decoupled Gibbs work (~free on a dispatch-bound accelerator; real cost on CPU, so leave off there). Template (temperature-linear) mode only. Probe statistics are computed on the sliced live prefix, but the probe RNG stream differs from an unpadded run, so discovered N can shift within its normal probe-to-probe variability.	False

Returns:

Type

Description

dict

dict with keys: n_chains: final recommended chain count betas: optimized schedule at that chain count Lambda: conservative (max) barrier estimate Lambda_raw: last raw estimate (may be lower than Lambda) target_acceptance: the target used converged_reason: "chain_count" | "lambda_stable" | "max_iters" barrier_identified: whether the final count's ladder round-tripped, so the reported Lambda is identified rather than a stalled-conveyor within-basin artifact (None if round trips were not tracked). history: list of per-probe dicts (each carries barrier_identified)

hamon.tune_exploration(key: jax.Array, ebm_factory: Callable | None = None, program_factory: Callable | None = None, init_states: Sequence[list] = (), clamp_state: list | None = None, initial_betas: jax.Array | None = None, start_steps: int = 1, max_steps: int = 64, rounds_per_probe: int = 400, n_tune_per_probe: int = 3, improve_tol: float = 0.05, time_rounds: int = 200, time_reps: int = 3, cost_model: bool = True, select_by: str = 'cost', target_efficiency: float = 0.9, *, fixed_schedule: jax.Array | None = None, ebm: AbstractEBM | None = None, program: BlockSamplingProgram | None = None, device: DeviceLike = 'auto') -> dict

Discover the local-exploration count gibbs_steps_per_round (n_expl).

n_expl is the only major NRPT knob hamon does not otherwise auto-tune (tune_chains sets N from Λ; tune_schedule sets the schedule). The objective maximized here is effective sample size per measured steady-state wall-second,

objective(n_expl) = ESS_median(cold chain) / (rounds · t_round(n_expl)),

where t_round is the per-round wall time measured on the target device after warm-up (so XLA compile is excluded). This is the honest endgame and, crucially, it self-calibrates to the hardware's real cost structure. The per-round cost is t_round = c₀ + n_expl·c_s: a fixed overhead c₀ (swap pass, energy recompute, host/kernel dispatch, scan bookkeeping) plus n_expl Gibbs sweeps at c_s each. The Syed et al. compute model assumes c₀ = 0 (cost ∝ n_expl), which makes n_expl=1 optimal; but cold-chain ESS grows sub-linearly in n_expl while real per-round cost grows less than linearly when c₀ is non-trivial — so on a dispatch-bound backend where c₀ ≳ 1.4·c_s the optimum shifts to n_expl > 1. Measuring t_round rather than assuming cost ∝ n_expl is what lets the search see this.

Why not the cheaper round-trip proxy τ_obs / n_expl? Empirically it under-picks n_expl: round trips count excursions but not how decorrelated each cold sample is, so doubling n_expl can sharply raise cold-chain ESS while barely moving τ_obs (the r≈0.81 ESS↔round-trip correlation of Syed et al. breaking down). The round-trip rate is still used as a robust gate and cross-check — efficiency = τ_obs/τ̄ is the ELE-violation meter, and the efficiency_limiter from :func:report_nrpt_diagnostics stops the search when a probe is schedule-limited (an unequalized ladder, where more local exploration cannot help). rt_per_compute and t_round are recorded per probe alongside the ESS-per-second objective.

See :func:tune_chains for the N analogue. The chain count is held fixed at len(initial_betas) — Λ (hence N) is a schedule invariant robust to n_expl, so the two searches decouple; run tune_chains first if N is unknown.

Each probe (1) gets a schedule, (2) measures ESS over a cold-chain trace via an :class:~hamon.NRPTStateObserver, and (3) times the steady-state round loop (warm-up absorbs the one-time compile; the median of time_reps runs of time_rounds rounds gives t_round). With fixed_schedule=None (default) the schedule is re-tuned per probe via :func:tune_schedule; pass fixed_schedule (a pre-tuned ladder) to reuse it and run each probe as a single :func:nrpt production call — much cheaper, and sound because the equi-acceptance schedule is invariant to n_expl. autotune uses this mode. Instead of ebm_factory / program_factory you may pass a single template ebm and program (temperature-linear mode).

Parameters:

Name	Type	Description	Default
key	Array	PRNG key.	required
ebm_factory / program_factory		per-chain factories, or use ebm / program template objects.	required
init_states	Sequence[list]	one initial block-state list per chain (fixed across probes).	()
clamp_state	list | None	clamped block states.	None
initial_betas	Array | None	the (fixed) β ladder; its length sets the chain count.	None
start_steps	int	smallest n_expl to try (≥ 1).	1
max_steps	int	largest n_expl before the search stops.	64
rounds_per_probe	int	production rounds per probe (and the tuning ceiling) — must be large enough for ESS and τ_obs to be low-variance.	400
n_tune_per_probe	int	schedule-tuning phases per probe.	3
improve_tol	float	minimum fractional objective gain to keep doubling (guards against chasing Monte-Carlo noise past the peak).	0.05
time_rounds	int	rounds per timed run when measuring t_round.	200
time_reps	int	number of timed runs to reduce over (noise control).	3
select_by	str	how to choose n_expl. "cost" (default) maximizes cold-chain ESS per wall-second (see cost_model) — the sample-quality objective; the pick depends on the machine's measured cost ratio, so it is best used as a one-time per-hardware calibration. "ele" instead picks the smallest count whose round-trip efficiency τ_obs/τ̂ reaches the ELE-adequacy knee — deterministic (no wall-clock) and the criterion the Syed et al. analysis prescribes, but it optimizes index-process mixing rather than cold-sample ESS, so it under-picks n_expl on a dispatch-bound accelerator where extra sweeps are nearly free; use it for index-efficiency or severe-ELE-violation regimes.	'cost'
target_efficiency	float	ELE-adequacy threshold for select_by="ele" — stop climbing once τ_obs/τ̂ reaches this (it also stops on a plateau or a schedule-limited verdict).	0.9
cost_model	bool	for select_by="cost" only — when True (and a fixed_schedule is given) fit t_round = c₀ + n_expl·c_s by least squares across the probes (each timing reuses the production executable, no separate observer=None compile) and take the argmax from that shared line; False times each probe independently (the flat objective then lets timing noise pick the count at random).	True
fixed_schedule	Array | None	a pre-tuned β ladder to reuse across all probes (each probe becomes one production run, no per-probe re-tuning). None (default) re-tunes per probe; autotune passes the ladder from tune_chains.	None
device	DeviceLike	where to run; resolved once and reused across probes. Timing is measured on this device, so the chosen n_expl is calibrated to it.	'auto'

Returns:

Type

Description

dict

dict with keys: gibbs_steps_per_round: the chosen n_expl. objective: ESS per measured wall-second at the chosen n_expl. ess_median / tau_observed / efficiency / rt_per_compute / t_round: at the choice (t_round in seconds). betas: the tuned schedule at the chosen n_expl. history: list of per-probe records (n_expl, objective, ess_median, tau_obs, rt_per_compute, t_round, efficiency, efficiency_limiter, betas).