API¶

torch_nf.exponential_families¶

Exponential family classes for implementing EFN.

class torch_nf.exponential_families.Dirichlet(D)¶

Bases: torch_nf.exponential_families.ExponentialFamily

KL(z, log_prob, eta)¶

T(z)¶

Sufficient statistic of the exponential family.

\(T(z) = \log(z)\)

\(\log(h(z)) = \sum_i \log(z_i)\)

The log base measure is concatenated onto the end of T(z).

Parameters: z (torch.tensor) – Random variable of the exponential family (M, N, D).
Returns: T(z) (M, N, D_eta).
Return type: torch.tensor

eta_to_mu(eta)¶

Obtain the mean parameterization (alpha) from the natural parameters (eta).

Parameters: eta (np.ndarray) – Mean parameterization (N, D_eta).
Returns: mu (N, D_mu).
Return type: np.ndarray

mu_to_eta(alpha)¶

Obtain the natural parameters (eta) from the mean parameterization (alpha).

Parameters: alpha (np.ndarray) – Mean parameterization (N, D_mu).
Returns: eta (N, D_eta).
Return type: np.ndarray

sample_eta(N=50, lb=0.5, ub=2.0)¶

Sample from prior distribution of the natural parameter eta.

\(\alpha_i \simi U[lb, ub]\)

A one is tacked onto the end of this natural parameter because there is a non-constant log base measure for the Dirichlet.

Parameters

N (int, optional) – Number of eta samples.
lb (float, optional) – Lower bound of prior.
ub (float, optional) – Upper bound of prior.

Returns

Samples of eta (N, D_eta).

Return type

np.ndarray

class torch_nf.exponential_families.ExponentialFamily(D, support_layer=None)¶

Bases: object

Exponential family distributions for training EFNs.

In this class, eta is the natural parameter and T calculates the sufficient statistics. Here, eta is the augmented natural parameter that has a 1 tacked onto the end when the log base measure is not a constant w.r.t. z. Likewise, T tacks on log(h(z)) when it is not a constant for the exponential family.

Parameters: D (int) – Dimensionality of random variable.

property D¶

T(z)¶

Sufficient statistic of the exponential family.

Parameters: z (torch.tensor) – Random variable of the exponential family (M, N, D).
Returns: T(z) (M, N, D_eta).
Return type: torch.tensor

eta_to_mu(eta)¶

Obtain the mean parameterization (mu) from the natural parameters (eta).

Parameters: eta (np.ndarray) – Mean parameterization (N, D_eta).
Returns: mu (N, D_mu).
Return type: np.ndarray

mu_to_eta(mu)¶

Obtain the natural parameters (eta) from the mean parameterization (mu).

Parameters: mu (np.ndarray) – Mean parameterization (N, D_mu).
Returns: eta (N, D_eta).
Return type: np.ndarray

sample_eta(N)¶

Sample from prior distribution of the natural parameter eta.

Parameters: N (int) – Number of eta samples.
Returns: Samples of eta (N, D_eta).
Return type: np.ndarray

property support_layer¶

class torch_nf.exponential_families.MVN(D)¶

Bases: torch_nf.exponential_families.ExponentialFamily

KL(z, log_prob, eta)¶

T(z)¶

Sufficient statistic of the exponential family.

\(T(z) = z, zz^{\top}\)

We eliminate redundancy in T(z) by vectorizing the upper triangular part of the second moment matrix.

Parameters: z (torch.tensor) – Random variable of the exponential family (M, N, D).
Returns: T(z) (M, N, D_eta).
Return type: torch.tensor

eta_to_mu(eta)¶

Obtain the mean parameterization (mu, Sigma) from the natural parameters (eta).

Parameters: eta (np.ndarray) – Mean parameterization (N, D_eta).
Returns: mu (N, D), Sigma (N, D, D).
Return type: np.ndarray, np.ndarray

mu_to_eta(mu, Sigma)¶

Obtain the natural parameters (eta) from the mean parameterization (mu, Sigma).

To avoid redundancy, we eliminate the copy of symmetric elements of the second moment natural parameter. We multiply such off diagonal elems by two in their vectorization in eta.

Parameters

mu (np.ndarray) – Mean (N, D)
Sigma (np.ndarray) – Covariance (N, D, D).

Returns

eta (N, D_eta).

Return type

np.ndarray

sample_eta(N=50, sigma_mu=1.0, iw_df_fac=5)¶

Sample from prior distribution of the natural parameter eta.

\(\mu_i \sim \mathcal{N}(0, \sigma_{\mu})\)

\(\Sigma \sim \mathcal{IW}(df=\text{iw_df_fac}*D, scale=\text{df_fac}*D*I)\)

Parameters

N (int, optional) – Number of eta samples.
sigma_mu (float, optional) – Standard deviation of Gaussian prior on \(\mu\).
iw_df_fac (int, optional) – Degree of freedom multiplier. df=iw_df_fac*D.

Returns

Samples of eta (N, D_eta).

Return type

np.ndarray

torch_nf.bijectors¶

class torch_nf.bijectors.Affine(D)¶

Bases: torch_nf.bijectors.Bijector

Affine bijector.

A scale and shift of each dimension. :param D: Dimensionality of the bijection. :type D: int

count_num_params()¶

Return the number of parameters for the bijector.

Returns: Number of parameters of the bijector.
Return type: int

forward_and_log_det(z, params)¶

Run the bijector forward on the input and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

inverse_and_log_det(z, params)¶

Run the bijector backwards and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

class torch_nf.bijectors.BatchNorm(D, momentum=0.1, eps=1e-05)¶

Bases: torch_nf.bijectors.Bijector

Batch Norm bijector.

For normalizing flows, it’s useful to have a BatchNorm layer that propagates its log-determinant jacobian. This Bijector keeps track of its most recently used normalization parameters, which can be invoked in the forward transform when use_last=True.

Parameters

D (int) – Dimensionality of the bijection.
momentum (float, optional.) – Momentum parameter of batch norm, default 0.1.
eps (float, optional.) – Eps parameter of batch norm, default 1e-5.

property eps¶

forward_and_log_det(z, use_last=False)¶

Batch norm forward and log determinant of the jacobian.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
use_last – Use previous mean and alpha of batch norm, default False.

get_last_alpha()¶

get_last_mean()¶

inverse_and_log_det(z)¶

Run the bijector backwards and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

property momentum¶

class torch_nf.bijectors.Bijector(D)¶

Bases: object

Base class for bijectors to be composed into normalizing flows.

Parameters: D (int) – Dimensionality of the bijection.

property D¶

count_num_params()¶

Return the number of parameters for the bijector.

Returns: Number of parameters of the bijector.
Return type: int

forward_and_log_det(z, params)¶

Run the bijector forward on the input and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

inverse_and_log_det(z, params)¶

Run the bijector backwards and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

class torch_nf.bijectors.MAF(D, num_layers, num_units, fwd_fac=True)¶

Bases: torch_nf.bijectors.Bijector

MAF bijector.

A masked autoregressive neural network parameterizes the affine bijection.

Parameters

D (int) – Dimensionality of the bijection.
num_layers (int) – Number of layers in the masked neural network.
num_units (int) – Number of hidden units per layer in network.

count_num_params()¶

Return the number of parameters for the bijector.

Returns: Number of parameters of the bijector.
Return type: int

forward_and_log_det(z, params)¶

Forward transform of MAF and log determinant of the jacobian.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

property fwd_fac¶

inverse_and_log_det(z, params)¶

Run the bijector backwards and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

property num_layers¶

property num_units¶

class torch_nf.bijectors.RealNVP(D, num_layers, num_units, transform_upper=True)¶

Bases: torch_nf.bijectors.Bijector

RealNVP bijector.

A fully connected neural network parameterizes the conditional affine transformation of half of the dimensions of the input z2 on the other half z1. The upper half is conditioned on the lower half if transform_upper=True, and the converse otherwise.

Parameters

D (int) – Dimensionality of the bijection.
num_layers (int) – Number of layers in the neural network for p(z2 | z1).
num_units (int) – Number of hidden units per layer in network for p(z2 | z1).
transform_upper (bool, optional) – z2 is up-half, and z1 is low-half, default True.

count_num_params()¶

Return the number of parameters for the bijector.

Returns: Number of parameters of the bijector.
Return type: int

forward_and_log_det(z, params)¶

Forward transform of the RealNVP and log determinant of the jacobian.

The shift and scale parameters of the upper (lower) half are outputs of two separate neural networks which take the lower (upper) half as input.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

inverse_and_log_det(z, params)¶

Run the bijector backwards and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

property num_layers¶

property num_units¶

property transform_upper¶

class torch_nf.bijectors.ToInterval(D, lb, ub)¶

Bases: torch_nf.bijectors.Bijector

Maps tensor in (M,N,D-1) to interval. :param D: Dimensionality of the bijection. :type D: int :param lb: Lower bound of interval. :type lb: float :param ub: Upper bound of interval. :type ub: float

forward_and_log_det(z)¶

Run the bijector forward on the input and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

inverse_and_log_det(z)¶

Run the bijector backwards and compute the log det jac.

The input z should have two batch dimensions. The first, of size M, should match the leading dimension of parameters params. The second batch dimension should be the number of samples per parameterization N. The earliest elements of the second dimension of z shall be used to parameterize this bijection.

Parameters

z (torch.tensor) – Input to the bijector (M, N, D).
params (torch.tensor) – Parameterization of the bijector (M, >|theta|).

property lb¶

property ub¶

class torch_nf.bijectors.ToSimplex(D)¶

Bases: torch_nf.bijectors.Bijector

Maps tensor in (M,N,D-1) to D-simplex.

Parameters: D (int) – Dimensionality of the bijection.

count_num_params()¶

Return the number of parameters for the bijector.

Returns: Number of parameters of the bijector.
Return type: int

forward_and_log_det(z)¶

Forward transform and log det jac of mapping to D-simplex.

Parameters: z (torch.tensor) – Input to the bijector (M, N, D).

torch_nf.bijectors.torch_atanh(x)¶