Masked coupling

[aseyboldt]

Add a second kind of normalizing flow, (I'm calling it masked-coupling for now).
This reduces compile time by a lot, but I'm not really sure how well it performs in terms of adaptability yet.

This can be activated with coupling_type="masked".

From some experimentation so far, it seems decent defaults might be something like

compiled = compiled.with_transform_adapt(
    num_layers=9,        # Number of layers in the normalizing flow
    nn_width=64,          # Neural networks with 32 hidden units
    num_diag_windows=6,   # Number of windows with a diagonal mass matrix intead of a flow
    verbose=True,         # Whether to print details about the adaptation process
    show_progress=True,   # Whether to show a progress bar for each optimization step
    coupling_type="masked",
    n_embed=None,  # New hyper parameter, default is 2 * nn_width right now
    n_deembed=None, # New hyper parameter, default is 2 * nn_width right now
)

trace = nutpie.sample(
    compiled,
    transform_adapt=True,
    seed=1,
    chains=1,
    cores=1,
    window_switch_freq=128,
)