amplfi.train.architectures module

class amplfi.train.architectures.flows.FlowArchitecture(num_params, embedding_net, embedding_weights=None, freeze_embedding=False)

Bases: Module

Base class for normalizing flow architectures that provides interface for interacting with embedding networks

build_flow()

Return type:

Flow

log_prob(x, context)

Wrapper around log_prob() from zuko.lazy.Flow object.

sample(n, context)

Wrapper around sample() from TransformedDistribution object.

class amplfi.train.architectures.flows.NSF(*args, transforms, hidden_features=(64, 64), passes=None, bins=8, randperm=False, residual=False, **kwargs)

Bases: FlowArchitecture

Light wrapper around the NSF flow from zuko library for compatibility with the FlowArchitecture interface.

See https://zuko.readthedocs.io/stable/api/zuko.flows.spline.html#zuko.flows.spline.NSF

Parameters:

• transforms (int) – Number of transformations in the flow

• hidden_features (Optional[Sequence[int]]) – Sequence of integers representing hidden units in the hyper network

• passes (Optional[int]) – Default of None corresponds to fully autoregressive flow. A value of 2 corresponds to coupling flow.

• bins (Optional[int]) – Number of bins in the spline

• randperm (Optional[bool]) – Whether to randomly permute features in between transformation layers

• residual (Optional[bool]) – Whether to use residual connections in the hyper network.

build_flow()

Return type:

Flow