tfep.nn.conditioners.made.MADE

class tfep.nn.conditioners.made.MADE(degrees_in: Sequence[int], degrees_out: Sequence[int], hidden_layers: int | Sequence[int] | Sequence[Sequence[int]] = 2, weight_norm: bool = True)[source]

Bases: Conditioner

An autoregressive layer implemented through masked affine layers.

This implements the Masked Autoregressive network for Distribution Estimation (MADE) [1], which is used in the Inverse/Masked Autoregressive Flow (IAF/MAF) [2]/[3] as a conditioner network. MADE is a dense layer, where the connections between nodes are partially masked to satisfy the autoregressive property. The mask is built based on the values of the degrees assigned to each node in the network (see [1]). Very briefly, an output node with degree i will depend only on the inputs assigned to a degree strictly less than i.

An advantage of using masks over the naive implementation of an autoregressive layer, which use a different neural network for each parameter of the affine transformation, is that it generates all the affine parameters in a single pass, with much less parameters to train, and can be parallelized trivially.

The current implementation supports arbitrary dependencies between input features so that this can be used as a conditioner to implement the full range between fully autoregressive and coupling layers flows.

Each layer is a MaskedLinear, which, in hidden layers, is followed by an ELU nonlinearity.

See also

generate_degrees(): Utility to generate sequences of degrees.

References

[1] Germain M, Gregor K, Murray I, Larochelle H. Made: Masked autoencoder: for distribution estimation. In International Conference on Machine Learning 2015 Jun 1 (pp. 881-889).
[2] Kingma DP, Salimans T, Jozefowicz R, Chen X, Sutskever I, Welling M.: Improved variational inference with inverse autoregressive flow. In Advances in neural information processing systems 2016 (pp. 4743-4751).
[3] Papamakarios G, Pavlakou T, Murray I. Masked autoregressive flow for: density estimation. In Advances in Neural Information Processing Systems 2017 (pp. 2338-2347).

Examples

A fully autoregressive MADE layer with 4 inputs and 8 outputs. Note that, in this example, because of how MADE masks are built from the degrees (see [1]), the first 2 outputs will not depend on any input (i.e., they will be trainable constant numbers). Similarly, the last input will not affect any output, and may be omitted from the model with no effect.

>>> made = MADE(
...     degrees_in=[0, 1, 2, 3],
...     degrees_out=[0, 0, 1, 1, 2, 2, 3, 3],
... )

The degrees do not have to be in any particular order, and multiple inputs can be assigned the same degree.

>>> made = MADE(
...     degrees_in=[1, 1, 0, 2],
...     degrees_out=[0, 1, 2, 3, 3, 2, 1, 0],
... )

It is possible to define “conditioning” inputs (in this case, the first 3). These are defined as features that affect all outputs and (consequently) for which no output is generated.

>>> made = MADE(
...     degrees_in=[-1, -1, -1, 0, 1, 2,],
...     degrees_out=[0, 1, 2, 3, 0, 1, 2, 3],
... )

The generate_degrees() utility function can be used to generate the degrees for several common scenarios. You can also control, the hidden layers.

>>> made = MADE(
...     degrees_in=generate_degrees(n_features=3, order='descending'),
...     degrees_out=generate_degrees(n_features=8, order='descending'),
...     hidden_layers=[
...         generate_degrees(n_features=6, max_value=2),
...         generate_degrees(n_features=8, max_value=2),
...     ],
... )

A coupling flow layer with 2 inputs, 4 outputs, and 3 hidden layers.

>>> made = MADE(
...     degrees_in=[-1, -1],
...     degrees_out=[0, 0, 0, 0],
...     hidden_layers=3,
... )

__init__(degrees_in: Sequence[int], degrees_out: Sequence[int], hidden_layers: int | Sequence[int] | Sequence[Sequence[int]] = 2, weight_norm: bool = True)[source]

Constructor.

Parameters:

degrees_in (Sequence[int]) – Shape: (n_inputs,). degrees_in[i] is the degree assigned to the i-th input.
degrees_out (Sequence[int]) – Shape: (n_outputs,). degrees_out[i] is the degree assigned to the i-th output.
hidden_layers (Union[int, Sequence[int], Sequence[Sequence[int]]], optional) – If an integer, this is the number of hidden layers. In this case, the number of nodes in each layer is set to max(n_inputs, ceil((n_inputs * n_outputs)**0.5)) where n_inputs is the number of input features that affect the output, and n_outputs is the number of output features.

If a sequence of integers, hidden_layers[l] is the number of nodes in the l-th hidden layer. The degrees of each node are assigned in a round-robin fashion by tiling degrees_in until the requested number of nodes is covered.

Otherwise, degrees_hidden[l][i] is the degree assigned to the i-th node of the l-th hidden layer.

Default is 2.
weight_norm (bool, optional) – If True, weight normalization is applied to the masked linear modules. Default is True.

Methods

`__init__`(degrees_in, degrees_out[, ...])	Constructor.
`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(x)	Compute the parameters for the transformer.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`modules`()	Return an iterator over all modules in the network.
`mtia`([device])	Move all model parameters and buffers to the MTIA.
`n_parameters`()	The total number of (unmasked) parameters.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post-hook to be run after module's `load_state_dict()` is called.
`register_load_state_dict_pre_hook`(hook)	Register a pre-hook to be run before module's `load_state_dict()` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_post_hook`(hook)	Register a post-hook for the `state_dict()` method.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`set_output`(output)	Implement `tfep.nn.flows.autoregressive.Conditioner.set_output()`.
`set_submodule`(target, module)	Set the submodule given by `target` if it exists, otherwise throw an error.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

Attributes

`T_destination`
`call_super_init`
`dimension_in`	Dimension of the input tensor.
`dimension_out`	Dimension of the output tensor.
`dimensions_hidden`	`(n_hidden_layers,)`.
`dump_patches`
`weight_norm`	`True` if weight norm is used, `False` otherwise.
`training`

property dimension_in: int: Dimension of the input tensor.

property dimension_out: int: Dimension of the output tensor.

property dimensions_hidden: Tensor

(n_hidden_layers,). dimensions_hidden[i] is the number of nodes in the i-th hidden layer.

Type:: Shape

forward(x)[source]

Compute the parameters for the transformer.

Parameters:: x (torch.Tensor) – Shape (batch_size, n_features). The input features.
Returns:: parameters – Shape (batch_size, n_parameters). The parameters for the transformer.
Return type:: torch.Tensor

n_parameters() → int[source]: The total number of (unmasked) parameters.

set_output(output: Tensor)[source]: Implement tfep.nn.flows.autoregressive.Conditioner.set_output().

property weight_norm

`True` if weight norm is used, `False` otherwise.

Type:: bool