DiscreteSampler¶

class DiscreteSampler : public madspace::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)

Public Functions

DiscreteSampler(const std::vector<std::size_t> &option_counts, const std::string &prefix = "", const std::vector<std::size_t> &dims_with_prior = {})¶

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.

inline const std::vector<std::size_t> &option_counts() const¶: Number of categories per discrete dimension.

inline const std::vector<std::string> &prob_names() const¶: Global names of the per-dimension probability vectors.

void initialize_globals(ContextPtr context) const¶: Register the probability globals on context.