DiscreteSampler

class madspace.DiscreteSampler(self: madspace._madspace_py.DiscreteSampler, option_counts: collections.abc.Sequence[SupportsInt], prefix: str = '', dims_with_prior: collections.abc.Sequence[SupportsInt] = [])

Bases: Mapping

Adaptive sampler for discrete choices.

Maps one uniform random number per discrete dimension to a category index, using a learned categorical distribution per dimension [1]. The probabilities are compute-graph globals named after prefix; call initialize_globals(context) once before use. A dimension listed in dims_with_prior is instead conditioned on a prior_i input.

batch is the leading batch dimension. d indexes the discrete dimensions.

Inputs

• random_i – float, shape (batch,) – one uniform number per dimension.

Conditions

• prior_d – float, shape (batch, option_counts[d]) – prior category weights. Present only for the dimensions in dims_with_prior.

Outputs

• index_i – int, shape (batch,) – the chosen category per dimension.

In addition every mapping returns a weight (float, shape (batch,)), the Jacobian of the transformation.

References

• [1] T. Heimel, O. Mattelaer, R. Winterhalder, “MadSpace”, https://arxiv.org/abs/2602.06895 (Sec. 3.2.2)

Parameters:

• option_counts – Number of categories for each discrete dimension.

• prefix – Prefix for the probability global names.

• dims_with_prior – Dimensions conditioned on a prior_i input rather than a learned global.

initialize_globals(self: madspace._madspace_py.DiscreteSampler, context: madspace._madspace_py.Context) → None

Register the probability globals on context.

option_counts(self: madspace._madspace_py.DiscreteSampler) → list[int]

Number of categories per discrete dimension.

prob_names(self: madspace._madspace_py.DiscreteSampler) → list[str]

Global names of the per-dimension probability vectors.