API Documentation

General

Main class (default PDF)

class jammy_flows.main.default.pdf(pdf_defs, flow_defs, options_overwrite={}, conditional_input_dim=None, amortization_mlp_dims='128', predict_log_normalization=False, join_poisson_and_pdf_description=False, hidden_mlp_dims_poisson='128', rank_of_mlp_mappings_poisson=0, amortization_mlp_use_custom_mode=False, amortization_mlp_ranks=0, amortization_mlp_highway_mode=0, amortize_everything=False, use_as_passthrough_instead_of_pdf=False, skip_mlp_initialization=False)
__init__(pdf_defs, flow_defs, options_overwrite={}, conditional_input_dim=None, amortization_mlp_dims='128', predict_log_normalization=False, join_poisson_and_pdf_description=False, hidden_mlp_dims_poisson='128', rank_of_mlp_mappings_poisson=0, amortization_mlp_use_custom_mode=False, amortization_mlp_ranks=0, amortization_mlp_highway_mode=0, amortize_everything=False, use_as_passthrough_instead_of_pdf=False, skip_mlp_initialization=False)

The main class of the project that defines a pytorch normalizing-flow PDF. The two main actions are evaluating the log-probability and sampling. Accessed via jammy_flows.pdf.

Parameters

  • pdf_defs (str) – String of characters describing the joint PDF structure: Sub-space structure is spearated by “+”. Example: “e2+s1+s2”, describes a joint PDF over a 2-dimensional euclidean space, a 1-sphere and a 2-sphere: a joint 5-dimensional PDF.

  • flow_defs (str) – A string, that describes how each conditional subfflow defined in “pdfs_defs” is structured in terms of normalizing-flow layers. Example: “gg+m+n” to describe a layer structure compatible with the other example “e2+s1+s2”. Two letters mean two consecutive applications of a certain flow layer, 3 letters three etc. Each layer holds their own parameters.

  • options_overwrite (dict) – Dictionary to overwrite default options of individual flow layers.

  • conditional_input_dim (None/int/list(int)) – Conditional input dimension if a conditional PDF. If a list is given, defines conditional input for each sub-pdf. None to define non-conditional PDF.

  • amortization_mlp_dims (str/list(str)) – Hidden structure of MLP for each sub-manifold.

  • predict_log_normalization (bool) – Predict log-mean of Poisson distribution

  • joint_poisson_and_pdf_description (bool) – If predicting log-mean, predict it together with other flow parameters if this is a conditional PDF. Only makes sense when conditional_input_dim is not None.

  • hidden_mlp_dims_poisson (str/list(str)) – If the log-mean is predicted by its own MLP, defines the hidden structure of the MLP.

  • amortization_mlp_use_custom_mode (bool) – Use custom AmortizableMLP class instead of default chained Linear layers.

  • rank_of_mlp_mappings_poisson (int) – Max rank of custom Poisson predictor MLP matrices.

  • amortization_mlp_highway_mode (int) – Connectivity mode for custom MLPs if used.

  • amortize_everything (bool) – Indicates whether all parameters, including the MLPs, should be amortized.

  • use_as_passthrough_instead_of_pdf (bool) – Indicates, whether the class acts as a PDF, or only as a flow mapping function of the overall autoregressive flow.

  • skip_mlp_initialization (bool) – Indicates, whether to skip MLP inits entirely. Can be used for custom MLP initialization.

all_layer_forward(x, log_det, data_summary, amortization_parameters=None, force_embedding_coordinates=False, force_intrinsic_coordinates=False)

Performs the autoregressive (IAF) forward normalizing-flow mapping of all sub-manifold flows.

Parameters

  • x (Tensor) – Target Input.

  • log_det (Tensor) – Input log-Det. Soon Deprecated.

  • data_summary (Tensor/None) – Holds summary information for the conditional PDF. Otherwise None.

  • amortization_parameters (Tensor/None) – Used to amortize the whole PDF. Otherwise None.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates in the output sample.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates in the output sample.

Returns

Tensor

Position in the target space after sampling.

Tensor

Log-Det factors of the forward mapping.

all_layer_inverse(x, log_det, data_summary, amortization_parameters=None, force_embedding_coordinates=False, force_intrinsic_coordinates=False)

Performs the autoregressive (IAF) backward normalizing-flow mapping of all sub-manifold flows.

Parameters

  • x (Tensor) – Target Input.

  • log_det (Tensor) – Input log-Det. Soon Deprecated.

  • data_summary (Tensor/None) – Holds summary information for the conditional PDF. Otherwise None.

  • amortization_parameters (Tensor/None) – Used to amortize the whole PDF. Otherwise None.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates in the input x for this inverse mapping.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates in the input x for this inverse mapping.

Returns

Tensor

Position in base space after inverse mapping.

Tensor

Log-Det factors of this inverse mapping.

approximate_coverage(target_x, conditional_input=None, amortization_parameters=None, force_embedding_coordinates=False, force_intrinsic_coordinates=False, num_percentile_points=100, sub_manifolds=[-1])

Calculates approximate coverage via the base distribution with the quantity 2*(log(p(0))-log(p_(z_base))) which should be chi^2 distributed for good coverage.

Parameters

  • target_x (Tensor) – Target positions to calculate coverage with. Must be of shape (B,D), where B = batch dimension.

  • conditional_input (Tensor/list(Tensor)/None) – Amortization input for conditional PDFs. If given, must be of shape (B,A), where A is the conditional input dimension defined in __init__. Can also be a list of tensors, one for each sub-PDF, if conditional_input_dim in __init__ is a list of ints.

  • amortization_parameters (Tensor/None) – If the PDF is fully amortized, defines all the parameters of the PDF. Must be of shape (B,T), where T is the total number of parameters of the PDF.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates in the input x.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates in the input x.

  • num_percentile_points (int) – At how many points along the chi2 do we want to compare true vs expected coverage?

  • sub_manifolds (list(int)) – Contains indices of sub-manifolds if coverage should be calculated for onditional PDF of the given sub-manifold. -1 stands for the total PDF and is the default.

Returns

Array of expected coverage probabilities. actual_coverage_probs (Numpy array): Array of actual coverage probabilities. actual_twice_logprob (Numpy array): Array of twice the log-probability difference at the base.

Return type

expected_coverage_probs (Numpy array)

count_parameters(verbose=False)

Counts parameters of the model. It does not matter, if all paramters are amortized or not, will always return the same.

Parameters

verbose (bool) – Prints out number of parameters. Differentiates amortization from non-amortization params.

Returns

int

Number of parameters (incl. amortization params).

coverage_and_or_pdf_scan(labels=None, conditional_input=None, amortization_parameters=None, coverage_num_percentile_points=100, sub_manifolds=[-1], exact_coverage_calculation=False, save_pdf_scan=False, calculate_MAP=False)

Calculates coverage (approximate) and possibly exact. Performs pdf scan for exact coverage and save scan if desired. The pdf scan and the exact coverage must be calculated in intrinsic coordinates (e.g. theta/phi for direction instead of dx/dy/dz).

Parameters

  • labels (Tensor) – Target positions to calculate coverage with. Must be of shape (B,D), where B = batch dimension.

  • conditional_input (Tensor/list(Tensor)/None) – Amortization input for conditional PDFs. If given, must be of shape (B,A), where A is the conditional input dimension defined in __init__. Can also be a list of tensors, one for each sub-PDF, if conditional_input_dim in __init__ is a list of ints.

  • amortization_parameters (Tensor/None) – If the PDF is fully amortized, defines all the parameters of the PDF. Must be of shape (B,T), where T is the total number of parameters of the PDF.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates in the input x.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates in the input x.

  • coverage_num_percentile_points (int) – At how many points along the chi2 do we want to compare observed vs expected coverage?

  • sub_manifolds (list(int)) – Contains indices of sub-manifolds if coverage should be calculated for onditional PDF of the given sub-manifold. -1 stands for the total PDF and is the default.

  • exact_coverage_calculation (bool) – Calculate exact coverage based on pdf scan?

  • save_pdf_scan (bool) – Save a pdf scan?

  • calculate_MAP (bool) – Calculate Maximum APosterior (MAP) coordinates based on pdf scan?

Returns

Dictionary of requested coverage and/or pdf scan values.

Return type

return_dict (dict)

entropy(sub_manifolds=[-1], conditional_input=None, force_embedding_coordinates=True, force_intrinsic_coordinates=False, samplesize=100, failsafe_crosscheck_tolerance=None, dtype=None, device=None)

Calculates entropy of the PDF.

Parameters

  • sub_manifolds (list(int)) – Contains indices of sub-manifolds if entropy should be calculated for marginal PDF of the given sub-manifold. -1 stands for the total PDF and is the default.

  • conditional_input (Tensor/list(Tensor)/None) – If passed defines the input to the PDF.

  • force_embedding_coordinates (bool) – Forces embedding coordinates in entropy calculation. Should always be true for correct manifold entropies.

  • force_intrinsic_coordinates (bool) – Forces intrinsic coordinates in entropy calculation. Should always be false for correct manifold entropies.

  • samplesize (int) – Samplesize to use for entropy approximation.

  • failsafe_crosscheck_tolerance (float / None) – If set, is used to crosscheck forward/bakckward pass compatability and resample if necessary. Has been introduced for the v flow in particular, so it should not be necessary for other flows.

  • dtype (torch dtype) – If given, uses this dtype. Otherwise uses dtype from parameters.

  • device (torch.device) – If given, uses this device. Otherwise uses device from parameters.

Returns

dict

Dictionary containing entropy for each index defined in parameter sub_manifolds. If -1 was given in sub_manifolds, the resulting entropy is stored under the total key.

entropy_iterative(sub_manifolds=[-1], conditional_input=None, force_embedding_coordinates=True, force_intrinsic_coordinates=False, samplesize=100, iterative_samplesize=10, max_iterative_batchsize=20, failsafe_crosscheck_tolerance=None, dtype=None, device=None, return_samples=False, verbose=False)

Calculates entropy of the PDF in an iterative manner. By iterating potentially both over target samples of later sub-pdfs, and over batch items, memory is saved and larger samplesizes (> a few 100 - 1000s) can be calculated. Probably only necessary for large sample sizes of later sub-pdfs, not of the first sub-pdf or the total PDF.

Parameters

  • sub_manifolds (list(int)) – Contains indices of sub-manifolds if entropy should be calculated for marginal PDF of the given sub-manifold. -1 stands for the total PDF and is the default.

  • conditional_input (Tensor/list(Tensor)/None) – If passed defines the input to the PDF.

  • force_embedding_coordinates (bool) – Forces embedding coordinates in entropy calculation. Should always be true for correct manifold entropies.

  • force_intrinsic_coordinates (bool) – Forces intrinsic coordinates in entropy calculation. Should always be false for correct manifold entropies.

  • samplesize (int) – Samplesize to use for entropy approximation.

  • iterative_samplesize (int) – Number of target PDF samples evaluated simultaneously. Must be a divisor of samplesize.

  • max_iterative_batchsize (int) – The max number of batch samples evaluated simultaneously.

  • failsafe_crosscheck_tolerance (float / None) – If set, is used to crosscheck forward/bakckward pass compatability and resample if necessary. Has been introduced for the v flow in particular, so it should not be necessary for other flows.

  • dtype (torch dtype) – If given, uses this dtype. Otherwise uses dtype from parameters.

  • device (torch.device) – If given, uses this device. Otherwise uses device from parameters.

  • return_samples (bool) – Return the samples that are generated to calculate the entropy? Samples are returned as B*num_samples X sample_dim, so the effective batch dimension is B*num_samples.

  • verbose (bool) – Adds some extra prints if given.

Returns

dict

Dictionary containing entropy for each index defined in parameter sub_manifolds. If -1 was given in sub_manifolds, the resulting entropy is stored under the total key.

Tensor

Only returned if return_samples is set to True. A tensor that contains the generated samples used to calculate the entropy.

forward(x, conditional_input=None, amortization_parameters=None, force_embedding_coordinates=False, force_intrinsic_coordinates=False)

Calculates log-probability at the target x. Also returns some other quantities that are calculated as a consequence.

Parameters

  • x (Tensor) – Target position to calculate log-probability at. Must be of shape (B,D), where B = batch dimension.

  • conditional_input (Tensor/list(Tensor)/None) – Amortization input for conditional PDFs. If given, must be of shape (B,A), where A is the conditional input dimension defined in __init__. Can also be a list of tensors, one for each sub-PDF, if conditional_input_dim in __init__ is a list of ints.

  • amortization_parameters (Tensor/None) – If the PDF is fully amortized, defines all the parameters of the PDF. Must be of shape (B,T), where T is the total number of parameters of the PDF.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates in the input x.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates in the input x.

Returns

Tensor

Log-probability, shape = (B,)

Tensor

Log-probability at base distribution, shape = (B,)

Tensor

Position at base distribution, shape = (B,D)

get_embedding_flags()

Returns embedding parameters flags (True/False) for different sub pdfs. An embedding flag is important for manifold sub dimensions. Set to True means the given sub-pdf flows expect the input to be in embedding space, while set to False the sub-pdf flows expect their input to be in intrinsic coordinates.

get_total_embedding_dim()

Returns embedding dimension of the overall PDF.

init_params(data=None, damping_factor=1000.0, mvn_min_max_sv_ratio=0.0001)

Initialize params of the normalizing flow such that the different sub flows play nice with each other and the starting distribution is a reasonable one. For the Gaussianization flow, data can be used to initilialize the starting distribution such that it roughly follows the data.

Parameters

  • data (None/Tensor) – If given a Tensor with target data, Gaussianization Flow subflows can make use of the distribution and initialize such that they follow the distribution.

  • damping_factor (float) – Weights in final matrices of amortization MLPs are divided by this factor (after already having been initialized) to dampen the impact of previous flow layers and conditional input in the autoregressive amortization structure.

log_mean_poisson(conditional_input=None, amortization_parameters=None)

Calculates log-mean Poisson prediction.

Parameters

  • conditional_input (Tensor/None) – Must be tensor of appropriate dimension if conditional PDF. Otherwise None.

  • amortization_parameters (Tensor/None) – Used to amortize the whole PDF. Otherwise None.

Returns

Tensor

log-lambda of Poisson, size (B,1)

marginal_moments(conditional_input=None, samplesize=50, iterative_samplesize=10, max_iterative_batchsize=20, mises_abs_precision=1e-07, calc_kl_diff_and_entropic_quantities=False, failsafe_crosscheck_tolerance=None, dtype=None, device=None, verbose=False, s2_entropy_scanning=False, s2_entropy_scan_nside=32)

Calculate the first and second central moments of the marginal distributions. For Euclidean manifolds it calculates a Gaussian approximation, for spherical distributions calculates a von-Mises approximation. Because these are the respective maximum entropy distributions, their entropy should always be larger than the original distribution. We can also calculate the kl divergence and cross entropy between the exact distribution and its approximation, by switching on the calc_kl_diff_and_entropic_quantities flag.

Parameters

  • conditional_input (Tensor/list(Tensor)/None) – If passed defines the input to the PDF.

  • samplesize (int) – Samplesize to use per event for moment approximation.

  • iterative_samplesize (int) – Number of target PDF samples evaluated simultaneously. Must be a divisor of samplesize.

  • max_iterative_batchsize (int) – The max number of batch samples evaluated simultaneously.

  • mises_abs_precision (float) – The absolute precision used break out the loop for the mises distribution second moment estimation.

  • calc_kl_diff_and_entropic_quantities (bool) – Flag that determines if we also calculate the KL divergence between the respective marginal distribution and its respective 2nd-order approximation. Also includes the cross entropy.

  • failsafe_crosscheck_tolerance (float / None) – If set, is used to crosscheck forward/bakckward pass compatability and resample if necessary. Has been introduced for the v flow in particular, so it should not be necessary for other flows.

  • dtype (torch dtype) – If given, uses this dtype. Otherwise uses dtype from parameters.

  • device (torch.device) – If given, uses this device. Otherwise uses device from parameters.

  • verbose (bool) – Some extra print statements on runtime.

  • s2_entropy_scanning (bool) – Use a healpix scan to determine entropy .. can be faster for certain s2 distributions.

Returns

dict

Dictionary containing moments and potentially entropies and KL-divergence between approximation and marginal distributions for each marginal distribution. The index indicates which marginal distribution. Means on the sphere are given in embedding coordinates and in angle coordinates. For example:

mean_0 = mean of first marginal in embedding coordinates mean_0_angles = mean of first marginal in angle coordinates varlike_0 = covariance (Gaussian/Euclidean) or concentration parameter (Fisher-von Mises/spherical) entropy_0 = entropy of first marginal distribution kl_diff_exact_approx_0 = KL divergence between exact distribution and 2nd-order approximation cross_entropy_0 = cross entropy between exact distribution and 2nd-order approximation … mean_1 = mean of second marginal … … entropy_total = total entropy (only total quantity calculated because it is essentially free)

sample(conditional_input=None, samplesize=1, seed=None, allow_gradients=False, amortization_parameters=None, force_embedding_coordinates=False, force_intrinsic_coordinates=False, failsafe_crosscheck_tolerance=None, dtype=None, device=None)

Samples from the (conditional) PDF.

Parameters

  • conditional_input (Tensor/list(Tensor)/None) – Tensor of shape B x D where B is the batch size and D the input space dimension if given. Can also be a list of tensors, which must share batch dimensionality. Else None.

  • samplesize (int) – Samplesize.

  • seed (None/int) –

  • allow_gradients (bool) – If False, does not propagate gradients and saves memory by not building the graph. Off by default, so has to be switched on for training.

  • amortization_parameters (Tensor/None) – Used to amortize the whole PDF. Otherwise None.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates for the sample.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates for the sample.

  • dtype (torch dtype) – Dtype and device are normally inferred by parameters or conditional input. If no parameters are part of the

  • device (torch.device) – If given, uses this device. Otherwise uses device from parameters.

Returns

Tensor

Sample in target space.

Tensor

Sample in base space.

Tensor

Log-pdf evaluation in target space

Tensor

Log-pdf evaluation in base space

set_embedding_flags(usement_flag, sub_pdf_index=None)

Resets the default parameter embedding structure. This defines how tensors are to be given to the PDF without extra flags like force_embedding_coordinates or force_intrinsic_coordinates.

Parameters

  • usement_Flag (bool) – Defines if selected sub manifold should switch to embedding coordinates (True) or intrinsic coordinates (False).

  • sub_pdf_index (int/None) – The index of of the manifold for which to set the flag. If None, sets flag for all.

transform_target_into_returnable_params(target)

Transforms an input tensor from default to embedding parametrization.

Parameters

target (Tensor) – Input tensor in current default parametrization.

Returns

Tensor

Output tensor in embedding parametrization.

transform_target_space(target, log_det=0, transform_from='default', transform_to='embedding')

Transform the destimation space tensor of the PDF as defined by transform_from and transform_to, which is the embedding space. (I.e. transform spherical angles into x/y/z pairs and so on when transform_to is embedding space.)

Parameters

  • target (Tensor) – Tensor to transform.

  • log_det (Tensor) – log_det Tensor (Soon Deprecated)

  • transform_from (str) – Coordinates to start with. One of “default”/”intrinsic”/”embedding”.

  • transform_to (str) – Coordinates to end with. One of “default”/”intrinsic”/”embedding”.

Returns

Tensor

Target in new coordinates.

Tensor

Any additional log-det added to input log-det.

Fully amortizable PDF

class jammy_flows.main.fully_amortized.fully_amortized_pdf(pdf_defs, flow_defs, options_overwrite={}, conditional_input_dim=None, inner_mlp_dims_sub_pdfs='128', inner_mlp_ranks=0, inner_mlp_highway_mode=1, amortization_mlp_dims='128', amortization_mlp_use_custom_mode=True, amortization_mlp_ranks=5, amortization_mlp_highway_mode=0, predict_log_normalization=False, skip_mlp_initialization=False)
__init__(pdf_defs, flow_defs, options_overwrite={}, conditional_input_dim=None, inner_mlp_dims_sub_pdfs='128', inner_mlp_ranks=0, inner_mlp_highway_mode=1, amortization_mlp_dims='128', amortization_mlp_use_custom_mode=True, amortization_mlp_ranks=5, amortization_mlp_highway_mode=0, predict_log_normalization=False, skip_mlp_initialization=False)

A fully amortized PDF, where in contrast to the standard autoregressive conditional PDF, also the whole autoregressive transformation, including MLPs, is amortized.

See documentation for an exact structural difference. Accessed via jammy_flows.fully_amortized_pdf.

Parameters

  • pdf_defs (str) – String of characters describing the joint PDF structure: Sub-space structure is spearated by “+”. Example: “e2+s1+s2”, describes a joint PDF over a 2-dimensional euclidean space, a 1-sphere and a 2-sphere: a joint 5-dimensional PDF.

  • flow_defs (str) – A string, that describes how each conditional subfflow defined in “pdfs_defs” is structured in terms of normalizing-flow layers. Example: “gg+m+n” to describe a layer structure compatible with the other example “e2+s1+s2”. Two letters mean two consecutive applications of a certain flow layer, 3 letters three etc. Each layer holds their own parameters.

  • options_overwrite (dict) – Dictionary to overwrite default options of individual flow layers.

  • conditional_input_dim (None/int) – Conditional input dimension if a conditional PDF.

  • inner_mlp_dims_sub_pdfs (str/list(str)) – Hidden structure of MLP for each sub-manifold.

  • inner_mlp_ranks (int/str) – The maximum rank allowed for the inner custom mlp mappings.

  • inner_mlp_highway_mode (int) – The highway mode used for the custom MLP. See amortizable_mlp class for more details.

  • amortization_mlp_dims (str/list(str)) – Hidden structure of the amortization MLP.

  • amortization_mlp_use_custom_mode (bool) – Indicates, whether custom mode should be used for all MLPs in the autoregressive linking process.

  • amortization_mlp_ranks (int/str) – The maximum rank allowed for the amortization MLP. See amortizable_mlp class for more details.

  • amortization_mlp_highway_mode (int) – The highway mode used for the amortization MLP. See amortizable_mlp class for more details.

  • predict_log_normalization (bool) – Predict log-mean of Poisson distribution.

  • skip_mlp_initialization (bool) – Indicates, whether to skip MLP inits entirely. Can be used for custom MLP initialization.

count_parameters(verbose=False)

Counts parameters of the model. It does not matter, if all paramters are amortized or not, will always return the same.

Parameters

verbose (bool) – Prints out number of parameters. Differentiates amortization from non-amortization params.

Returns

int

Number of parameters (incl. amortization params).

forward(x, conditional_input=None, force_embedding_coordinates=False, force_intrinsic_coordinates=False)

Calculates log-probability at the target x. Also returns some other quantities that are calculated as a consequence.

Parameters

  • x (Tensor) – Target position to calculate log-probability at. Must be of shape (B,D), where B = batch dimension.

  • conditional_input (Tensor/None) – Amortization input for conditional PDFs. If given, must be of shape (B,A), where A is the conditional input dimension defined in __init__.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates in the input x.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates in the input x.

Returns

Tensor

Log-probability, shape = (B,)

Tensor

Log-probability at base distribution, shape = (B,)

Tensor

Position at base distribution, shape = (B,D)

sample(conditional_input=None, samplesize=1, seed=None, allow_gradients=False, force_embedding_coordinates=False, force_intrinsic_coordinates=False)

Samples from the (conditional) PDF.

Parameters

  • conditional_input (Tensor/None) – Of shape N x D where N is the batch size and D the input space dimension if given. Else None.

  • samplesize (int) – Samplesize.

  • seed (None/int) –

  • allow_gradients (bool) – If False, does not propagate gradients and saves memory by not building the graph. Off by default, so has to be switched on for training.

  • force_embedding_coordinates (bool) – Enforces embedding coordinates for the sample.

  • force_intrinsic_coordinates (bool) – Enforces intrinsic coordinates for the sample.

Returns

Tensor

Sample in target space.

Tensor

Sample in base space.

Tensor

Log-pdf evaluation in target space

Tensor

Log-pdf evaluation in base space

Other options and functions

jammy_flows.flow_options.obtain_default_options(flow_abbrevation)

Obtains a dictionary with default options for a given flow. For a complete list of possible options, have a look in flow_options.py.

Parameters

flow_abbreviation (str) – The character specifying a particular flow layer.

Returns

dict

Dictionary containing the default options.

jammy_flows.extra_functions.log_one_plus_exp_x_to_a_minus_1(x, a)

Calculates the logarithm of ((1+exp(x))**a - 1) / ((1+exp(x))**a)

Parameters

  • x (Tensor) – Size (B,D)

  • a (Tensor) – Size (B,1)

Returns

Tensor

Result

Amortizable MLP

class jammy_flows.amortizable_mlp.AmortizableMLP(input_dim, hidden_dims, output_dim, highway_mode=0, low_rank_approximations=0, nonlinearity='tanh', use_permanent_parameters=True, svd_mode='smart', precise_mlp_structure={})
__init__(input_dim, hidden_dims, output_dim, highway_mode=0, low_rank_approximations=0, nonlinearity='tanh', use_permanent_parameters=True, svd_mode='smart', precise_mlp_structure={})

MLP class that allows for amortization of all of its parameters. Supports low-rank approximations to weight matrices and 5 different settings of connectivity (“highway modes”).

Parameters

  • dim (output) – Input dimension of the overall module.

  • hidden_dims (str/int/list(int)) – Defines hidden dimension structure of all MLPs, depending on highway_mode how they are used. Hidden dimensionality is given by integers, separated by dash. I.e. 64-64.

  • dim – Output dimension of the overall module.

  • highway_mode (int) –

    Defines connectivity structure.

    0: Default MLP with a given number of hidden layers defined by hidden_dims.
    1: Like mode 0 but with an extra skip connection using a linear matrix connecting input and output.
    2: Like mode 1 but multiple one-hidden-layer MLPs that all connect from input to output instead of the single MLP. Each hidden layer of those MLPs is defined by hidden_dims.
    3: Like mode 2 but the 2nd to nth MLP connects from the previous output to the next output, instead of input to output. Additionally, there are skip connections between all stages.
    4: Like mode 3, but the 2nd to nth MLP connections from a union of [input,previous output] to next output. Additionally, there are skip connections between all stages.

  • low_rank_approximation (int/list(int)) – Defines low-rank approximation of all matrices. If 0 or smaller, full rank is assumed.

  • nonlinearity (str) – Nonlinearity used.

  • use_permanent_parameters (bool) – If False, is used in amortization mode, i.e. all parameters are input in foward as extra_inputs.

  • svd_mode (str) – One of [“naive”, “smart”]. Smart takes a full normal matrix if the rank is its max rank, then a naive low-rank approximation is less parameter efficient. Naive just always makes a SVD based low-rank approximation, which can be less efficient sometimes.

  • precise_mlp_structure (dict) – Allows to pass a dict with similar structure as sub_mlp_structures (see code), in order to have more customization, i.e. per-matrix varying low-ranks etc.

forward(i, extra_inputs=None)

Applies the AmortizableMLP to target.

Parameters

  • i (Tensor) – Input tensor of shape (B, D)

  • extra_inputs (Tensor/None) – If given, amortizes all parameters of the AmortizableMLP.

Returns

Tensor

Output of the AmortizableMLP.

initialize_uvbs(fix_total=None, fix_final_bias=None, prev_damping_factor=1000.0)

Initialize the parameters of the MLPs. Either all parameters are initialized similar to default procedures in pytorch (fix_total and fix_final_bias None), or all parameters are handed over (fix_total given), or the final bias vector is handed over (fix_final_bias given).

Parameters

  • fix_total (None/Tensor) – If not None, has to be of size of all AmortizableMLP params to initialize all params of all MLPs.

  • fix_final_bias (None/Tensor) – If not None, has to be of size of the target dimension (final bias) to set the final bias vector.

  • prev_damping_factor (float) – Used to dampen previous weights/biases of init tensor when fix_final_bias is given.

obtain_default_init_tensor(fix_final_bias=None, prev_damping_factor=1000.0)

Obtains a tensor that fits the shape of the AmortizableMLP parrameters according to standard kaiman init rules taken from pytorch.

Parameters

  • fix_final_bias (Tensor/None) – If given, the final bias is fixed to a specific output Tensor given by fix_final_bias. If None, standard initialization applies.

  • prev_damping_factor (float) – If fix_final_bias is set, determines a damping factor that applies to all weights/biases before final bias, in order to reduce dependency on them.

Returns

Tensor

The final initialization tensor that is desired for the AmortizableMLP.

Numerical inversion

jammy_flows.layers.bisection_n_newton.inverse_bisection_n_newton(func, grad_func, target_arg, *args, min_boundary=-100000.0, max_boundary=100000.0, num_bisection_iter=25, num_newton_iter=30, newton_tolerance=1e-14, verbose=0)

Performs bisection and Newton iterations simulataneously in each 1-d subdimension in a given batch.

Parameters

  • func (function) – The function to find the inverse of.

  • grad_func (function) – The gradient function of the function.

  • target_arg (float Tensor) – The argument at which the inverse functon should be evaluated. Tensor of size (B,D) where B is the batchsize, and D the dimension.

  • *args (list) – Any extra arguments passed to func.

  • min_boundary (float) – Minimum boundary for Bisection.

  • max_boundary (float) – Maximum boundary for bisection.

  • num_bisection_iter (int) – Number of bisection iterations.

  • num_newton_iter (int) – Number of Newton iterations.

Returns

Tensor

The inverse of the function func in each sub-dimension in each batch item.

jammy_flows.layers.bisection_n_newton.inverse_bisection_n_newton_joint_func_and_grad(func, joint_func, target_arg, *args, min_boundary=-100000.0, max_boundary=100000.0, num_bisection_iter=25, num_newton_iter=30, newton_tolerance=1e-14, verbose=0)

Performs bisection and Newton iterations simulataneously in each 1-d subdimension in a given batch.

Parameters

  • func (function) – The function to find the inverse of.

  • joint_func (function) – A function that calculates the function value and its derivative simultaneously.

  • target_arg (float Tensor) – The argument at which the inverse functon should be evaluated. Tensor of size (B,D) where B is the batchsize, and D the dimension.

  • *args (list) – Any extra arguments passed to func.

  • min_boundary (float) – Minimum boundary for Bisection.

  • max_boundary (float) – Maximum boundary for bisection.

  • num_bisection_iter (int) – Number of bisection iterations.

  • num_newton_iter (int) – Number of Newton iterations.

Returns

Tensor

The inverse of the function func in each sub-dimension in each batch item.

jammy_flows.layers.bisection_n_newton.inverse_bisection_n_newton_sphere(combined_func, find_tangent_func, basic_exponential_map_func, target_arg, *args, num_newton_iter=25)

Performs Newton iterations on the sphere over 1-dimensional potential functions via Exponential maps to find the inverse of a given exponential map on the sphere. In initial tests it was found that a very precise application requires at least 40-50 ierations, even though one is already pretty close after 10 iterations. Distributing the points randomly on the sphere doesn’t really help, so every point is initialized just at (0,0,1). In order for this to work properly, the function has to be globally diffeomorphic, so the exponential map has to follow the conditions outlined in https://arxiv.org/abs/0906.0874 (Sei 2009).

Parameters

  • combined_func (function) – A function that returns the (x,y,z) unit vector and its jacobian.

  • find_tangent_func (function) – A funtion to calculate the tangent at a given point along a certain direction.

  • target_arg (float Tensor) – The argument at which the inverse functon should be evaluated. Tensor of size (B,D) where B is the batchsize, and D the dimension.

  • *args (list) – Any extra arguments passed to func.

  • num_newton_iter (int) – Number of Newton iterations.

Returns

Tensor

The inverse of the exponential map.

jammy_flows.layers.bisection_n_newton.inverse_bisection_n_newton_sphere_fast(combined_func, find_tangent_func, basic_exponential_map_func, target_arg, *args, num_newton_iter=25)

Performs Newton iterations on the sphere over 1-dimensional potential functions via Exponential maps to find the inverse of a given exponential map on the sphere. In initial tests it was found that a very precise application requires at least 40-50 ierations, even though one is already pretty close after 10 iterations. Distributing the points randomly on the sphere doesn’t really help, so every point is initialized just at (0,0,1). In order for this to work properly, the function has to be globally diffeomorphic, so the exponential map has to follow the conditions outlined in https://arxiv.org/abs/0906.0874 (Sei 2009).

Parameters

  • combined_func (function) – A function that returns the (x,y,z) unit vector and its jacobian.

  • find_tangent_func (function) – A funtion to calculate the tangent at a given point along a certain direction.

  • target_arg (float Tensor) – The argument at which the inverse functon should be evaluated. Tensor of size (B,D) where B is the batchsize, and D the dimension.

  • *args (list) – Any extra arguments passed to func.

  • num_newton_iter (int) – Number of Newton iterations.

Returns

Tensor

The inverse of the exponential map.

Helper Functions

Layers

class jammy_flows.layers.layer_base.layer_base(dimension=1, always_parametrize_in_embedding_space=0)
__init__(dimension=1, always_parametrize_in_embedding_space=0)

Base class for all flow layers.

Parameters

  • dimension (int) – The intrinsic dimension of the flow.

  • always_parametrize_in_embedding_space (int) – Required for subclasses to make decisions for embedding transformations.

_embedding_conditional_return(x)

Transforms a vector into embedded space. Is flow specific and must be implemented by subclass.

_embedding_conditional_return_num()

Returns the number of embedding space dimensions for this flow. Must be implemented by subclass.

flow_mapping(input, extra_inputs=None)

Implements the flow forward mapping, which is required for sampling. Must be overwritten by subclass.

get_desired_init_parameters()

Returns desired init parameters for given flow. Should be overwritten by subclass.

init_params(params)

Initializes permanent params. Can only be invoked when PDF is non-conditional. Must be overwritten by subclass.

inv_flow_mapping(input, extra_inputs=None)

Implements the flow inverse mapping, which is required for evaluation. Must be overwritten by subclass.

obtain_layer_param_structure(param_dict, extra_inputs=None, previous_x=None, extra_prefix='')

Useful for debugging purposes. Must be overwritten by subclass.

Parameters

param_dict (dict) – Is handed over from jammy_flows.pdf class and filled in this function with current values via param_dict[par_name]=*par_value*.

transform_target_space(x, log_det=0.0, trafo_from='default', trafo_to='embedding')

Transforms the target vector from a given parametrization to another parametrization. Must be implemented by subclass.

Euclidean flow layers

Euclidean Base

class jammy_flows.layers.euclidean.euclidean_base.euclidean_base(dimension=1, use_permanent_parameters=False, model_offset=0)
__init__(dimension=1, use_permanent_parameters=False, model_offset=0)

Base class for all Euclidean flow layers. Inherits from layer_base.

Parameters

model_offset (int) – If set, models an extra offset for the Euclidean flow. Will be done automatically by jammy_flows pdf class. Only makes sense for the last flow layer in a sequence of Euclidean flows.

Identity layer (“x”)

class jammy_flows.layers.euclidean.euclidean_do_nothing.euclidean_do_nothing(dimension, use_permanent_parameters=True, add_offset=0)
__init__(dimension, use_permanent_parameters=True, add_offset=0)

Identitiy transformation. Symbol “x”

Multivariate Normal (“t”)

class jammy_flows.layers.euclidean.multivariate_normal.mvn_block(dimension, cov_type='full', use_permanent_parameters=False, model_offset=0, width_smooth_saturation=1, lower_bound_for_widths=0.01, upper_bound_for_widths=100, softplus_for_width=0, clamp_widths=0)
__init__(dimension, cov_type='full', use_permanent_parameters=False, model_offset=0, width_smooth_saturation=1, lower_bound_for_widths=0.01, upper_bound_for_widths=100, softplus_for_width=0, clamp_widths=0)

Affine flow (Multivariate Normal Distribution) - Symbol “t”

Parameters

  • cov_type (str) – One of [“full”, “diagonal”, “diagonal_symmetric”, “identity”]. “diagonal_symmetric” means one shared covariance parameter for all dimensions.

  • width_smooth_saturation (int) – If set, uses smooth function also for lower end to describe log-width of variance-like parameters.

  • lower_bound_for_widths (float) – Lower bound for the variance-like parameters.

  • upper_bound_for_widths (float) – Upper bound for the variance-like parameters.

  • softplus_for_width (int) – If set, uses softplus instead of exponential to enforce positivity on variance-like parameters.

  • clamp_widths (int) – If set, clamps widths.

Gaussianization flow (“g”)

class jammy_flows.layers.euclidean.gaussianization_flow.gf_block(dimension, nonlinear_stretch_type='classic', num_kde=5, num_householder_iter=-1, use_permanent_parameters=False, fit_normalization=0, inverse_function_type='inormal_partly_precise', model_offset=0, softplus_for_width=0, width_smooth_saturation=1, lower_bound_for_widths=0.01, upper_bound_for_widths=100, lower_bound_for_norms=1, upper_bound_for_norms=10, center_mean=0, clamp_widths=0, regulate_normalization=0, add_skewness=0, rotation_mode='householder')
__init__(dimension, nonlinear_stretch_type='classic', num_kde=5, num_householder_iter=-1, use_permanent_parameters=False, fit_normalization=0, inverse_function_type='inormal_partly_precise', model_offset=0, softplus_for_width=0, width_smooth_saturation=1, lower_bound_for_widths=0.01, upper_bound_for_widths=100, lower_bound_for_norms=1, upper_bound_for_norms=10, center_mean=0, clamp_widths=0, regulate_normalization=0, add_skewness=0, rotation_mode='householder')

Gaussianization flow: Symbol “g”

Implements https://arxiv.org/abs/2003.01941 as a modified version of the official implementation in https://github.com/chenlin9/Gaussianization_Flows. Fixes numerical issues with bisection inversion due to more efficient newton iterations, which also makes samples differentiable + fixes numerical issues with inverse CDF parametrization. Also allows for amortization (the original implementation is meant only for non-conditional PDFs). Further adds multiple rotation parametrizations, logistic skewness option, and an rq-spline option for the non-linear mapping.

Parameters

  • dimension (int) – Dimension of the PDF. Set by jammy_flows.pdf.

  • nonlinear_stretch_type (str) – One of [“classic”, “rq_splines”]. Default is classic, but also allows to replace non-linear mapping with rational quadratic splines.

  • num_kde (int) – Number of KDE Kernels in each one-dimensional transformation.

  • num_householder_iter (int) – if <=0, no householder transformation is performed. If positive, it defines the number of parameters in householder transformations.

  • fit_normalization (int) – If set, also fits the normalization.

  • inverse_function_type (str) – One of [“partly_crude”, “partly_precise”, “full_pade”, “isgimoid”]. “Partly_crude” is implemented in the original repository (https://github.com/chenlin9/Gaussianization_Flows), but has numerical issues. The option “isigmoid” “is an easier inverse CDF that looses exact CDF identification but gains some speed and is numerically stable. It is recommended to use “isigmoid” or “partly_precise”.

  • model_offset (int) – Set by jammy_flows.pdf.

  • softplus_for_width (int) – If set, uses soft_plus to predict width. Otherwhise exponential function.

  • width_smooth_saturaiton (int) – If set, saturates the width symmetrically at the lower and upper end.

  • lower_bound_for_widths (float) – Determines the lower width bound for each KDE element.

  • upper_bound_for_widths (float) – Determines the upper width bound for each KDE element.

  • lower_bound_for_norms (float) – Determines the lower norm bound for each KDE element.

  • upper_bound_for_norms (float) – Determines the upper norm bound for each KDE element.

  • clamp_widths (int) – If set, clamps widths above threshold.

  • regulate_normalization (int) – If set, regulates with a lower bound of 1 and upper bound of 100 (fixed) similar to width regulation.

  • add_skewness (int) – If set, adds skewness to the basis functions in order to make flow more flexible.

  • rotation_mode (str) –

    One of different rotation parametrizations:

    - householder: Default implementation from the original paper.
    - angles: Parametrized via “Givens rotation”
    - triangular_combination: Parametrization of an upper and lower triangular matrix.
    - cayley: “Cayley representation” of rotations.

Polynomial stretch flow (“p”)

class jammy_flows.layers.euclidean.polynomial_stretch_flow.psf_block(dimension, num_transforms=3, num_householder_iter=-1, use_permanent_parameters=False, model_offset=0, exact_mode=True)
__init__(dimension, num_transforms=3, num_householder_iter=-1, use_permanent_parameters=False, model_offset=0, exact_mode=True)

Polynomial stretch flow - Symbol: “p”

A polynomial flow that uses abs(x) to make arbitrary polynomials work as flow-mappings. Has interesting structure but does not seem to work very well in practice.

Spherical flow layers

jammy_flows.layers.spheres.sphere_base.return_safe_costheta(x, safety_margin=None)

Restricts the angle to not hit 0 or pi exactly.

class jammy_flows.layers.spheres.sphere_base.sphere_base(dimension=1, euclidean_to_sphere_as_first=True, use_permanent_parameters=False, rotation_mode='householder', add_rotation=False, higher_order_cylinder_parametrization=False)
__init__(dimension=1, euclidean_to_sphere_as_first=True, use_permanent_parameters=False, rotation_mode='householder', add_rotation=False, higher_order_cylinder_parametrization=False)

Base class for all spherical flow layers. Inherits from layer_base.

add_rotation (int): If set, adds an additional rotation as an additional “flow”. rotation_mode (str): One of [“angles”, “householder”]. Angles involves Givens Rotations, and householder householder rotations.

Identity layer (“y”)

class jammy_flows.layers.spheres.spherical_do_nothing.spherical_do_nothing(dimension, euclidean_to_sphere_as_first=False, use_permanent_parameters=True, add_rotation=0)
__init__(dimension, euclidean_to_sphere_as_first=False, use_permanent_parameters=True, add_rotation=0)

Spherical identity mapping: Symbol “y”

Moebius 1-d circle flow (“m”)

class jammy_flows.layers.spheres.moebius_1d.moebius(dimension=1, euclidean_to_sphere_as_first=True, add_rotation=0, natural_direction=0, use_permanent_parameters=False, use_moebius_xyz_parametrization=True, num_basis_functions=5)
__init__(dimension=1, euclidean_to_sphere_as_first=True, add_rotation=0, natural_direction=0, use_permanent_parameters=False, use_moebius_xyz_parametrization=True, num_basis_functions=5)

Moebius transformations on the circle. Symbol “m”

Follows https://arxiv.org/abs/2002.02428. Parametrization seems to be slightly unstable sometimes.

Parameters

  • natural_direction (int) – If set to 0, log-probability is faster than sampling. Otherwise reversed.

  • use_moebius_xyz_parametreization (bool) – Two different paramerizations.

  • num_basis_functions (int) – Number of moebius basis functions.

Circular 1-d spline flow (“o”)

class jammy_flows.layers.spheres.splines_1d.spline_1d(dimension=1, euclidean_to_sphere_as_first=True, add_rotation=1, natural_direction=0, use_permanent_parameters=False, num_basis_functions=10)
__init__(dimension=1, euclidean_to_sphere_as_first=True, add_rotation=1, natural_direction=0, use_permanent_parameters=False, num_basis_functions=10)

A circular spline variant. Symbol: “o”

Follows https://arxiv.org/abs/2002.02428. Still experimental and has not really been tested.

Segmented 2-d sphere flow (“n”)

class jammy_flows.layers.spheres.segmented_sphere_nd.segmented_sphere_nd(dimension, euclidean_to_sphere_as_first=True, add_rotation=1, rotation_mode='householder', use_permanent_parameters=False, use_moebius_xyz_parametrization=True, num_basis_functions=5, zenith_type_layers='ggt', max_rank=20, hidden_dims='64', subspace_mapping='logistic', higher_order_cylinder_parametrization=False, highway_mode=0)
__init__(dimension, euclidean_to_sphere_as_first=True, add_rotation=1, rotation_mode='householder', use_permanent_parameters=False, use_moebius_xyz_parametrization=True, num_basis_functions=5, zenith_type_layers='ggt', max_rank=20, hidden_dims='64', subspace_mapping='logistic', higher_order_cylinder_parametrization=False, highway_mode=0)

A segmented conditional flow that models a flow on the torus and then maps to the sphere. Symbol “n”

Follows https://arxiv.org/abs/2002.02428. Seems to be stable when non-conditional, but can be unstable in certain configurations in conditional mode. Slight modification, in that instead of the torus one can also choose an infinite torus, whose z-axis is modeled via a Euclidean flow instead of an interval flow.

Parameters

  • rotation_mode (str) – One of [“angles, “householder”]

  • use_moebius_xyz_parametreization (bool) – Two different parametrizations for the Moebius part.

  • num_basis_functions (int) – Number of Moebius basis functions.

  • zenith_type_layers (str) – Either “g” or multiple “gg” (Gaussianization flows) for infinite torus parametrization, or “r” or multiple “rr” (Rational-Quadratic Spline) for torus parametrization.

  • max_rank (int) – Max rank of the MLP that maps from azimuth to zenith dimension.

  • hidden_dims (str/int/list(int)) – Hidden dims of MLP that maps from azimuth to zenith dimension.

  • higher_order_cylinder_parametrization (bool) – Slightly different parametrization which allows some simplifications in intermedate steps. Default False (makes it hard to combine spherical flows).

num_parameter_list = None

self.flow_dict[“g”] = dict() self.flow_dict[“g”][“module”] = gf_block self.flow_dict[“g”][“type”] = “e” self.flow_dict[“g”][“kwargs”] = dict() self.flow_dict[“g”][“kwargs”][“fit_normalization”] = 1 self.flow_dict[“g”][“kwargs”][“use_permanent_parameters”]=0 self.flow_dict[“g”][“kwargs”][“num_householder_iter”] = -1 self.flow_dict[“g”][“kwargs”][“num_kde”] = 10 self.flow_dict[“g”][“kwargs”][“inverse_function_type”] = “isigmoid”

self.flow_dict[“g”][“kwargs”][“softplus_for_width”]=0 # use softplus instead of exp to transform log_width -> width self.flow_dict[“g”][“kwargs”][“upper_bound_for_widths”]=100 # define an upper bound for the value of widths.. -1 = no upper bound self.flow_dict[“g”][“kwargs”][“lower_bound_for_widths”]=0.01 # define a lower bound for the value of widths self.flow_dict[“g”][“kwargs”][“clamp_widths”]=0 self.flow_dict[“g”][“kwargs”][“width_smooth_saturation”]=1 #

self.flow_dict[“p”] = dict() self.flow_dict[“p”][“module”] = psf_block self.flow_dict[“p”][“type”] = “e” self.flow_dict[“p”][“kwargs”] = dict() self.flow_dict[“p”][“kwargs”][“use_permanent_parameters”]=0 self.flow_dict[“p”][“kwargs”][“num_householder_iter”] = -1 self.flow_dict[“p”][“kwargs”][“num_transforms”] = 3

self.flow_dict[“x”] = dict() self.flow_dict[“x”][“module”] = euclidean_do_nothing self.flow_dict[“x”][“type”] = “e” self.flow_dict[“x”][“kwargs”]=dict()

### interval layers - for these layers, the input (0-pi) only has to be transformed to 0-1 via a cos-transformation

self.flow_dict[“r”] = dict() self.flow_dict[“r”][“module”] = rational_quadratic_spline self.flow_dict[“r”][“type”] = “i” self.flow_dict[“r”][“kwargs”] = dict() self.flow_dict[“r”][“kwargs”][“use_permanent_parameters”]=0 self.flow_dict[“r”][“kwargs”][“num_basis_elements”]=5 self.flow_dict[“r”][“kwargs”][“low_boundary”] = 0 self.flow_dict[“r”][“kwargs”][“high_boundary”] = 1.0

self.flow_dict[“z”] = dict() self.flow_dict[“z”][“module”] = interval_do_nothing self.flow_dict[“z”][“type”] = “i” self.flow_dict[“z”][“kwargs”]=dict()

Exponential map 2-d sphere flow (“v”)

class jammy_flows.layers.spheres.exponential_map_s2.exponential_map_s2(dimension, euclidean_to_sphere_as_first=False, use_permanent_parameters=False, exp_map_type='linear', natural_direction=0, num_components=10, add_rotation=0, max_num_newton_iter=1000, mean_parametrization='old')
__init__(dimension, euclidean_to_sphere_as_first=False, use_permanent_parameters=False, exp_map_type='linear', natural_direction=0, num_components=10, add_rotation=0, max_num_newton_iter=1000, mean_parametrization='old')

Uses the spherical exponential map. Symbol: “v”

Uses linear and quadratic potential as described in https://arxiv.org/abs/0906.0874, and exponential potential as described in https://arxiv.org/abs/2002.02428. Additionally added a spline-based potential.

Parameters

  • exp_map_type (str) – Defines the potential of the exponential map. One of [“linear”, “quadratic”, “exponential”, “splines”].

  • natural_direction (int) –

  • num_components (int) – How many components to sum over in the exponential map.

  • add_rotation (int) – Add a rotation after the main flow?.

  • num_newton_iter (int) – Number of newton iterations.

  • mean_parametrization (str) – “old” (x,y,z, directly), or “householder”, which desribes each mean by a unit vec and rotates them to new positions.

Chart-based continuous 2-d sphere flow (“c”)

class jammy_flows.layers.spheres.cnf_sphere_charts.cnf_sphere_charts(dimension, euclidean_to_sphere_as_first=False, use_permanent_parameters=False, cnf_network_hidden_dims='64-64', cnf_network_rank=0, cnf_network_highway_mode=1, num_charts=6, solver='rk4', atol=1e-07, rtol=1e-07, step_size=0.03125)
__init__(dimension, euclidean_to_sphere_as_first=False, use_permanent_parameters=False, cnf_network_hidden_dims='64-64', cnf_network_rank=0, cnf_network_highway_mode=1, num_charts=6, solver='rk4', atol=1e-07, rtol=1e-07, step_size=0.03125)

Continuous manifold normalizing flow - Symbol: “c”

Code implementation of “Neural manifold ordinary differential equations” (https://arxiv.org/abs/2006.10254). Mostly adapted from https://github.com/CUAI/Neural-Manifold-Ordinary-Differential-Equations.

Parameters

  • cnf_network_hidden_dims (str/int/list(int)) – Hidden dim structure of MLP.

  • cnf_network_rank – (int, list(int)): Low rank structure of MLP.

  • cnf_network_highway_mode (int) – Highway connectivity mode of MLP.

  • num_charts (int) – How many charts to use.

  • solver (str) – One of [“rk4”, “dopri5”, “dopri8”, “bosh3”, “fehlberg2”, “adaptive_heun”, “euler”, “midpoint”].

  • atol (float) – Absolute tolerance. (Used for adaptive solvers like dopri)

  • rtol (float) – Relative tolerance. (Used for adaptive solvers like dopri)

  • step_size (float) – Step size for fixed step solvers (like rk4, euler).

Interval flow layers

class jammy_flows.layers.intervals.interval_base.interval_base(dimension=1, euclidean_to_interval_as_first=0, use_permanent_parameters=False, low_boundary=0.0, high_boundary=1.0)
__init__(dimension=1, euclidean_to_interval_as_first=0, use_permanent_parameters=False, low_boundary=0.0, high_boundary=1.0)

Identity layer (“z”)

class jammy_flows.layers.intervals.interval_do_nothing.interval_do_nothing(dimension, euclidean_to_interval_as_first=0, use_permanent_parameters=False, low_boundary=0.0, high_boundary=1.0)
__init__(dimension, euclidean_to_interval_as_first=0, use_permanent_parameters=False, low_boundary=0.0, high_boundary=1.0)

Interval identity mapping: Symbol “z”

Rational-quadratic spline flow (“r”)

class jammy_flows.layers.intervals.rational_quadratic_spline.rational_quadratic_spline(dimension, num_basis_functions=10, euclidean_to_interval_as_first=0, use_permanent_parameters=0, low_boundary=0, high_boundary=1.0, min_width=0.0001, min_height=0.0001, min_derivative=0.0001, fix_boundary_derivatives=-1.0, smooth_second_derivative=0, restrict_max_min_width_height_ratio=-1.0)
__init__(dimension, num_basis_functions=10, euclidean_to_interval_as_first=0, use_permanent_parameters=0, low_boundary=0, high_boundary=1.0, min_width=0.0001, min_height=0.0001, min_derivative=0.0001, fix_boundary_derivatives=-1.0, smooth_second_derivative=0, restrict_max_min_width_height_ratio=-1.0)

Rational-quadartic spline layer: Symbol “r”

Also known as Neural spline flows - https://arxiv.org/abs/1906.04032. Adapted code from pyro implementation.

Parameters

  • num_basis_functions (int) – Number of spline elements.

  • low_boundary (float) – Lower boundary of the inverval.

  • high_boundary (float) – Higher boundary of the interval.

  • min_width (float) – Minimum width of a spline element.

  • min_height (float) – Minimum height of a spline element.

  • min_derivative (float) – Minimum derivative of a spline element. Only used if smooth_second_derivative is 0.

  • fix_boundary_derivatives (float) – If larger than 0, determines the value of the boundary derivatives

  • smooth_second_derivative (int) – Determines if second derivatives should be smooth. Only works for 2 basis function currently. Ignores min_derivative if positive.

  • restrict_max_min_width_height_ratio (float) – Maximum relative between maximum/miniumum of widths/heights.. negative will not constrain this quantity.

Simplex flow layers

class jammy_flows.layers.simplex.simplex_base.simplex_base(dimension=1, use_permanent_parameters=False, always_parametrize_in_embedding_space=0, project_from_gauss_to_simplex=0)
__init__(dimension=1, use_permanent_parameters=False, always_parametrize_in_embedding_space=0, project_from_gauss_to_simplex=0)

Base class for all simplex flow layers. Inherits from layer_base.

Iterative interval simplex flow (“w”)

class jammy_flows.layers.simplex.inner_loop_simplex.inner_loop_simplex(dimension, use_permanent_parameters=False, always_parametrize_in_embedding_space=0, project_from_gauss_to_simplex=0)
__init__(dimension, use_permanent_parameters=False, always_parametrize_in_embedding_space=0, project_from_gauss_to_simplex=0)

Iterative simplex flow. Symbol: “w”

Construction suggested in https://arxiv.org/abs/2008.05456.