DiscreteFlow

class DiscreteFlow : public madspace::Mapping

Autoregressive adaptive sampler for discrete choices.

Like DiscreteSampler, but the categorical distribution for each discrete dimension is predicted by an MLP subnetwork from the earlier choices and an optional conditioning input [1]. The subnetwork weights are compute-graph globals named after prefix; call initialize_globals(context) once before use.

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

Inputs

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

Conditions

  • condition – float, shape (batch, condition_dim) – conditioning input. Present only when condition_dim is nonzero.

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

Outputs

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

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

References

Public Functions

DiscreteFlow(const std::vector<std::size_t> &option_counts, const std::string &prefix = "", const std::vector<std::size_t> &dims_with_prior = {}, std::size_t condition_dim = 0, std::size_t subnet_hidden_dim = 32, std::size_t subnet_layers = 3, MLP::Activation subnet_activation = MLP::leaky_relu)
Parameters:

  • option_counts – Number of categories for each discrete dimension.

  • prefix – Prefix for the trainable global names.

  • dims_with_prior – Dimensions conditioned on a prior_d input.

  • condition_dim – Width of the conditioning input; 0 for none.

  • subnet_hidden_dim – Hidden width of the MLP subnetworks.

  • subnet_layers – Number of layers in the subnetworks.

  • subnet_activation – Activation of the subnetworks.

inline const std::vector<std::size_t> &option_counts() const

Number of categories per discrete dimension.

inline std::size_t condition_dim() const

Width of the conditioning input.

void initialize_globals(ContextPtr context, std::optional<std::uint64_t> seed = std::nullopt) const

Register the flow’s trainable parameters on context.

Parameters:

  • context – Context receiving the globals.

  • seed – Seed for the random initialization. If unset (the default), the initialization is non-deterministic.