API Documentation

General

Main class (default PDF)

Fully amortizable PDF

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.