---
id: zuko
version: "1.6.0"
license: MIT
license_treatment: permissive
maintenance: active
---

# zuko — Normalizing flows in PyTorch

License: permissive · Maintenance: active · Downloads: 76.3K/mo

## 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 above — 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

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

- License: MIT (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 76.3K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

normalizing flows pytorch, conditional probability distributions, generative modeling deep learning, neural density estimation, flow-based generative models, pytorch probabilistic inference, invertible transformations neural networks, probabilistic-modeling, generative-models, variational-inference

[View on SkillFed](https://skillfed.io/packages/zuko) · [View on PyPI](https://pypi.org/project/zuko/)