How antsx/antsnormalizingflows works
1. Repository shape
- This is a Python/PyTorch library, not an application service.
- The production code lives under src/antsnormflows, with a classic layered layout: core model composition, flow transforms, distributions, neural nets, sampling, and utility helpers.
- Surrounding it are examples/notebooks, docs, tests, and CI workflows.
2. Primary abstraction: NormalizingFlow
- The central architecture is in src/antsnormflows/core.py.
- A normalizing-flow model is built by composing:
- a base distribution,
- an ordered list of flow transforms.
- Data moves bidirectionally:
- training / density evaluation: observed samples are transformed back toward the base latent space while accumulating log-determinant terms,
- generation / sampling: latent samples are drawn from the base and pushed forward through the flow chain.
3. Flow subsystem
- src/antsnormflows/flows contains the transformation layers that make up the model.
- There are several families:
- Affine flows (affine/): coupling, autoregressive, Glow-style blocks.
- Neural spline flows (neural_spline/): coupling/autoregressive spline transforms and wrappers.
- Classical flows: planar, radial, residual, periodic, stochastic.
- Structural transforms: mixing, normalization, reshape.
- These modules are the main “graph” in the architecture diagram: they are chained together and operate on tensors plus log-determinants.
4. Conditioning and parameterization
- Flow layers typically do not store full transforms directly; they use small neural networks to predict transform parameters.
- That logic is shared across flow families and lives in src/antsnormflows/nets.
- In practice, a conditioner network (for example an MLP or MADE) feeds parameters into a coupling or autoregressive flow block.
5. Neural network building blocks
- src/antsnormflows/nets provides reusable parameter nets:
- mlp.py, cnn.py, made.py, resnet.py, lipschitz.py.
- These are support components, not a separate model family.
- Their role is to be embedded inside flows, encoder/decoder distributions, and any other module that needs learned parameter maps.
6. Distribution subsystem
- src/antsnormflows/distributions defines the probabilistic side of the library.
- It includes:
- base distributions and priors,
- target distributions,
- encoder/decoder distributions,
- proposal machinery for Metropolis-Hastings-style sampling,
- a linear interpolation helper.
- This layer is where the library connects flows to probabilistic modeling use cases, including latent-variable setups and target-density approximations.
7. Sampling subsystem
- src/antsnormflows/sampling/hais.py is a specialized sampling path separate from normal flow evaluation.
- It supports advanced annealed/importance-style sampling workflows rather than the core forward/inverse flow API.
- In architecture terms, this is an auxiliary inference/sampling subsystem that can consume the same distributions and transforms but serves a different purpose than standard training.
8. Shared utilities and low-level math
- src/antsnormflows/utils contains support code for:
- masks,
- spline math,
- neural-network helpers,
- preprocessing,
- evaluation,
- optimization helpers.
- src/antsnormflows/transforms.py appears to provide general transform utilities shared across multiple subpackages.
- These modules are implementation support rather than top-level conceptual subsystems.
9. Usage patterns reflected in examples
- examples/ shows the intended composition model: instantiate a base distribution, build a list of flow blocks, then wrap them in NormalizingFlow.
- The examples cover multiple model families: Real NVP, Glow, spline flows, planar/radial/residual flows, VAEs, conditional and image-oriented flows.
- This indicates the library is designed as a component toolkit for experimentation, not a single fixed pipeline.
10. Documentation and package delivery
- Documentation is built with MkDocs (mkdocs.yml, .github/workflows/mkdocs.yaml), and there is a generated PDF in docs/.
- Packaging is standard Python packaging (pyproject.toml, MANIFEST.in, requirements files), with the source-layout package under src/.
- The GitHub Actions workflows run tests, build docs, and produce the PDF artifact.
11. Test structure mirrors the package
- tests/ closely mirrors src/antsnormflows/, which makes the architecture explicit:
- core tests,
- distributions tests,
- flows tests by flow family,
- nets tests,
- utility tests.
- That mirroring is a strong indicator of the main boundaries in the codebase.
12. High-level dependency flow for a diagram
- The most important directional relationships are:
- NormalizingFlow depends on distributions + flows.
- flows depend on nets and utils.
- distributions may also depend on nets/utils for learned components.
- sampling builds on the probabilistic core.
- Everything is implemented in PyTorch tensors and modules, so the runtime boundary is the PyTorch execution model rather than any external service.
Main entry points: utils Utilities · examples Examples · mlp.py Parameter nets · core.py NormalizingFlow
Components 19
Library code
Primary PyTorch package under src/antsnormflows.
- NormalizingFlow · Central flow model that composes a base distribution with ordered transforms and handles forward/inverse density paths.
src/antsnormflows/core.py - Base dist · Base and prior distributions that seed latent-space sampling and density evaluation.
src/antsnormflows/distributions/base.py - Encoder/decoder · Latent-variable distribution components used for conditional inference and generative modeling.
src/antsnormflows/distributions/encoder.py - Target dist · Target-distribution helpers for approximation and probabilistic modeling workflows.
src/antsnormflows/distributions/target.py - Proposal tools · Metropolis-Hastings style proposal machinery for advanced sampling flows.
src/antsnormflows/distributions/mh_proposal.py - Flow base · Shared flow abstractions and base contracts for invertible tensor transforms.
src/antsnormflows/flows/base.py - Affine flows · Coupling, autoregressive, and Glow-style affine transforms.
src/antsnormflows/flows/affine/__init__.py - Spline flows · Neural spline coupling and autoregressive transforms plus wrappers.
src/antsnormflows/flows/neural_spline/__init__.py - Classical flows · Planar, radial, residual, periodic, and stochastic flow transforms.
src/antsnormflows/flows/planar.py - Structure ops · Mixing, normalization, and reshape transforms that organize flow topology and tensor shape.
src/antsnormflows/flows/mixing.py - HAIS sampling · Annealed or importance-style sampling workflow separate from the standard flow API.
src/antsnormflows/sampling/hais.py
Support code
Reusable helpers and low-level math used across the library.
- Parameter nets · Reusable MLP, CNN, MADE, ResNet, and Lipschitz networks that parameterize flow and distribution blocks.
src/antsnormflows/nets/mlp.py - Utilities · Masks, spline math, preprocessing, evaluation, neural-network helpers, and optimization support.
src/antsnormflows/utils - Transforms API · General transform utilities shared across subpackages.
src/antsnormflows/transforms.py
Examples and docs
Notebooks, scripts, and documentation showing intended composition and use.
Tests and CI
Mirrored tests and automation for validation and release.
- Core tests · Tests for model composition and core behavior.
tests/core_test.py - Flow tests · Mirrored tests for affine, spline, classical, and structural flow families.
tests/flows - Dist tests · Mirrored tests for base, encoder, decoder, prior, and target distribution logic.
tests/distributions
Connections 26
- NormalizingFlow → Base dist: uses · NormalizingFlow is built around a base distribution for latent sampling and density evaluation.
- NormalizingFlow → Flow base: composes · The model holds an ordered chain of flow transforms.
- Affine flows → Parameter nets: conditioned by · Affine coupling and autoregressive layers use learned parameter networks.
- Spline flows → Parameter nets: conditioned by · Spline flow blocks use neural parameterizers for knot and scale parameters.
- Classical flows → Parameter nets: sometimes uses · Classical flows may use learned subnetworks for residual or stochastic variants.
- Structure ops → Utilities: depends on · Structural transforms rely on masking and tensor helpers.
- Spline flows → Utilities: depends on · Spline transforms rely on spline math and helper utilities.
- Parameter nets → Utilities: built with · Reusable neural builders rely on shared helper code.
- Encoder/decoder → Parameter nets: uses · Encoder and decoder distributions are parameterized by learned networks.
- Encoder/decoder → Utilities: uses · Distribution helpers share preprocessing and helper utilities.
- Target dist → Utilities: uses · Target-density helpers depend on shared support routines.
- Proposal tools → Base dist: wraps · Proposal machinery operates in the same probabilistic layer as base and prior distributions.
- HAIS sampling → Base dist: consumes · HAIS sampling uses the probabilistic core as its input distribution layer.
- HAIS sampling → Flow base: steps through · Advanced sampling can traverse flow transforms as part of annealed execution.
- Transforms API → Utilities: shared math · General transform utilities are implemented on top of shared helpers.
- Examples → NormalizingFlow: demonstrates · Examples show how to compose a base distribution and flow stack into NormalizingFlow.
- Examples → Affine flows: demonstrates · Examples cover Real NVP and Glow-style affine flows.
- Examples → Spline flows: demonstrates · Examples cover neural spline flow variants.
- Examples → Classical flows: demonstrates · Examples cover planar, radial, residual, and related flows.
- Docs → Examples: documents · Documentation explains the library through notebooks and usage patterns.
- Core tests → NormalizingFlow: validates · Core tests cover model composition and forward/inverse behavior.
- Flow tests → Affine flows: validates · Flow tests mirror affine transform families.
- Flow tests → Spline flows: validates · Flow tests mirror spline transform families.
- Flow tests → Classical flows: validates · Flow tests mirror classical transforms.
- Dist tests → Base dist: validates · Distribution tests mirror base and prior components.
- Dist tests → Encoder/decoder: validates · Distribution tests mirror encoder and decoder components.