zuko

What it is and what it does

Zuko bridges a gap in PyTorch's probabilistic modeling: the built-in Distribution and Transform classes cannot be sent to GPU or treated as trainable modules. Zuko wraps these concepts in LazyDistribution and LazyTransform—modules whose forward pass returns a distribution or transformation—enabling conditional probability modeling and normalizing flows to work seamlessly within PyTorch's training pipeline.

The package provides a suite of flow architectures (NSF, RealNVP, MAF, NICE, and others) that can be instantiated with configurable hyperparameters and trained end-to-end. It also supports custom flows built from individual components like masked autoregressive transforms and rotation transforms, making it suitable for both standard use cases and research-driven extensions.

Use it for:

• Train a neural spline flow to model complex conditional distributions for density estimation or variational inference tasks.
• Build custom generative models by composing flow components like masked autoregressive transforms and base distributions.
• Implement probabilistic inference pipelines where conditional distributions must be parameterized by neural networks.
• Experiment with different normalizing flow architectures (RealNVP, MAF, NSF, etc.) for research or production density modeling.
• Sample from learned conditional distributions p(x|c) after training a flow on paired data.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Zuko implements normalizing flows in PyTorch as trainable neural modules, enabling conditional probability distributions and transformations that can be sent to GPU and optimized with standard PyTorch workflows.

Yes. Zuko fills a genuine need in PyTorch's probabilistic modeling ecosystem—making normalizing flows trainable and GPU-compatible. It is actively maintained, has no known vulnerabilities, low install friction, and a permissive MIT license. Install it if you need conditional probability distributions or flow-based generative models in PyTorch.

Install

zuko on PyPI

pip install zuko

uv add zuko

poetry add zuko

Installing zuko

Before you install

Low friction: pure Python wheel with only numpy and torch as runtime dependencies. Active maintenance with a release 157 days ago and ongoing repository activity.

License in practice

MIT license permits unrestricted use, modification, and distribution in commercial and private projects without warranty.

Quickstart

pip install zuko

import torch
import zuko

flow = zuko.flows.NSF(3, 5, transforms=3, hidden_features=[128] * 3)
optimizer = torch.optim.Adam(flow.parameters(), lr=1e-3)

for x, c in trainset:
    loss = -flow(c).log_prob(x)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Requires Python 3.10 or later and PyTorch installed.

Verify before relying

• Whether the package supports GPU acceleration beyond PyTorch's standard .to() mechanism.
• Performance characteristics and scalability limits for large-dimensional flows.
• Availability and completeness of tutorials for each flow type listed in the documentation.

Package facts

Evidence: zuko-1.6.0-py3-none-any.whl

Keywords: torch, normalizing flows, probability, distribution, generative, deep learning

Tags

More Artificial Intelligence packages

Further reading