core.nn.nsf.FlowWrapper

core.nn.nsf.FlowWrapper(flow, embedding_net=None)

This class wraps the neural spline flow. It is required for multiple reasons. (i) some embedding networks take tuples as input, which is not supported by the nflows package. (ii) paralellization across multiple GPUs requires a forward method, but the relevant flow method for training is log_prob.

Attributes

Name Description
embedding_net
flow

Methods

Name Description
forward
log_prob
sample
sample_and_log_prob

forward

core.nn.nsf.FlowWrapper.forward(y, *x)

log_prob

core.nn.nsf.FlowWrapper.log_prob(y, *x)

sample

core.nn.nsf.FlowWrapper.sample(*x, num_samples=1)

sample_and_log_prob

core.nn.nsf.FlowWrapper.sample_and_log_prob(*x, num_samples=1)