DiscreteFlow

class madspace.DiscreteFlow(self: madspace._madspace_py.DiscreteFlow, option_counts: collections.abc.Sequence[typing.SupportsInt], prefix: str = '', dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = [], condition_dim: typing.SupportsInt = 0, subnet_hidden_dim: typing.SupportsInt = 32, subnet_layers: typing.SupportsInt = 3, subnet_activation: madspace._madspace_py.MLP.Activation = <Activation.leaky_relu: 1>)

Bases: 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

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.

condition_dim(self: madspace._madspace_py.DiscreteFlow) → int

Width of the conditioning input.

initialize_globals(self: madspace._madspace_py.DiscreteFlow, context: madspace._madspace_py.Context, seed: SupportsInt | None = None) → None

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.

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

Number of categories per discrete dimension.