Package: zuko
Zuko - Normalizing flows in PyTorchZuko is a Python package that implements normalizing flows inPyTorch. It relies as much as possible on distributions and transformations already provided by PyTorch. Unfortunately, theDistributionandTransformclasses oftorchare not sub-classes oftorch.nn.Module, which means you cannot send their internal tensors to GPU with.to('cuda')or retrieve their parameters with.parameters(). Worse, the concepts of conditional distribution and transformation, which are essential for probabilistic inference, are impossible to express.To solve these problems,zukodefines two concepts: theLazyDistributionandLazyTransform, which are any modules whose forward pass returns aDistributionorTransform, respectively. Because the creation of the actual distribution/transformation is delayed, an eventual condition can be easily taken into account. This design enables lazy distributions, including normalizing flows, to act like distributions while retaining features inherent to modules, such as trainable parameters. It also makes the implementations easy to understand and extend.In theAvatarcartoon,Zukois a powerful firebender 🔥AcknowledgementsZuko takes significant inspiration fromnflowsandStefan Webb's work inPyroandFlowTorch.InstallationThezukopackage is available onPyPI, which means it is installable viapip.pip install zukoAlternatively, if you need the latest features, you can install it from the repository.pip install git+https://github.com/probabilists/zukoGetting startedNormalizing flows are provided in thezuko.flowsmodule. To build one, supply the number of sample and context features as well as the transformations' hyperparameters. Then, feeding a context $c$ to the flow returns a conditional distribution $p(x | c)$ which can be evaluated and sampled from.importtorchimportzuko# Neural spline flow (NSF) with 3 sample features and 5 context featuresflow=zuko.flows.NSF(3,5,transforms=3,hidden_features=[128]*3)# Train to maximize the log-likelihoodoptimizer=torch.optim.Adam(flow.parameters(),lr=1e-3)forx,cintrainset:loss=-flow(c).log_prob(x)# -log p(x | c)loss=loss.mean()optimizer.zero_grad()loss.backward()optimizer.step()# Sample 64 points x ~ p(x | c*)x=flow(c_star).sample((64,))Alternatively, flows can be built as customFlowobjects.fromzuko.flowsimportFlow,MaskedAutoregressiveTransform,Unconditionalfromzuko.distributionsimportDiagNormalfromzuko.transformsimportRotationTransformflow=Flow(transform=[MaskedAutoregressiveTransform(3,5,hidden_features=(64,64)),Unconditional(RotationTransform,torch.randn(3,3)),MaskedAutoregressiveTransform(3,5,hidden_features=(64,64)),],base=Unconditional(DiagNormal,torch.zeros(3),torch.ones(3),buffer=True,),)For more information, check out the documentation and tutorials atzuko.readthedocs.io.Available flowsClassYearReferenceGMM-Gaussian Mixture ModelNICE2014Non-linear Independent Components EstimationMAF2017Masked Autoregressive Flow for Density EstimationNSF2019Neural Spline FlowsNCSF2020Normalizing Flows on Tori and SpheresSOSPF2019Sum-of-Squares Polynomial FlowNAF2018Neural Autoregressive FlowsUNAF2019Unconstrained Monotonic Neural NetworksCNF2018Neural Ordinary Differential EquationsGF2020Gaussianization FlowsBPF2020Bernstein-Polynomial Normalizing FlowsContributingIf you have a question, an issue or would like to contribute, please read ourcontributing guidelines.