# antsx/antsnormalizingflows architecture

> Architecture diagram of the GitHub repository antsx/antsnormalizingflows, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-13). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/antsx/antsnormalizingflows
- Repository: https://github.com/antsx/antsnormalizingflows

## Overview

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: [`src/antsnormflows/utils`](https://github.com/antsx/antsnormalizingflows/tree/main/src/antsnormflows/utils) (Utilities), [`examples`](https://github.com/antsx/antsnormalizingflows/tree/main/examples) (Examples), [`src/antsnormflows/nets/mlp.py`](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/nets/mlp.py) (Parameter nets), [`src/antsnormflows/core.py`](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/core.py) (NormalizingFlow)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g_src["Library code"]
  node_core["NormalizingFlow<br/>model orchestrator<br/>[core.py]"]
  node_dist_base[("Base dist<br/>distribution core<br/>[base.py]")]
  node_dist_cond[("Encoder/decoder<br/>learned distributions<br/>[encoder.py]")]
  node_dist_target[("Target dist<br/>target density<br/>[target.py]")]
  node_dist_mh[("Proposal tools<br/>mh proposal<br/>[mh_proposal.py]")]
  node_flows_base["Flow base<br/>flow interface<br/>[base.py]"]
  node_flows_affine["Affine flows<br/>flow family<br/>[__init__.py]"]
  node_flows_spline["Spline flows<br/>flow family<br/>[__init__.py]"]
  node_flows_classic["Classical flows<br/>flow family<br/>[planar.py]"]
  node_flows_struct["Structure ops<br/>tensor transforms<br/>[mixing.py]"]
  node_sampling["HAIS sampling<br/>auxiliary sampler<br/>[hais.py]"]
end

subgraph group_g_support["Support code"]
  node_nets["Parameter nets<br/>neural builders<br/>[mlp.py]"]
  node_utils["Utilities<br/>shared helpers"]
  node_transforms["Transforms API<br/>shared transform utils<br/>[transforms.py]"]
end

subgraph group_g_usage["Examples and docs"]
  node_examples["Examples<br/>usage demos"]
  node_docs["Docs<br/>documentation<br/>[index.md]"]
end

subgraph group_g_quality["Tests and CI"]
  node_tests_core["Core tests<br/>test suite<br/>[core_test.py]"]
  node_tests_flows["Flow tests<br/>test suite"]
  node_tests_dist["Dist tests<br/>test suite"]
end

node_core -->|"uses"| node_dist_base
node_core -->|"composes"| node_flows_base
node_flows_affine -->|"conditioned by"| node_nets
node_flows_spline -->|"conditioned by"| node_nets
node_flows_classic -->|"sometimes uses"| node_nets
node_flows_struct -->|"depends on"| node_utils
node_flows_spline -->|"depends on"| node_utils
node_nets -->|"built with"| node_utils
node_dist_cond -->|"uses"| node_nets
node_dist_cond -->|"uses"| node_utils
node_dist_target -->|"uses"| node_utils
node_dist_mh -->|"wraps"| node_dist_base
node_sampling -->|"consumes"| node_dist_base
node_sampling -->|"steps through"| node_flows_base
node_transforms -->|"shared math"| node_utils
node_examples -->|"demonstrates"| node_core
node_examples -->|"demonstrates"| node_flows_affine
node_examples -->|"demonstrates"| node_flows_spline
node_examples -->|"demonstrates"| node_flows_classic
node_docs -->|"documents"| node_examples
node_tests_core -->|"validates"| node_core
node_tests_flows -->|"validates"| node_flows_affine
node_tests_flows -->|"validates"| node_flows_spline
node_tests_flows -->|"validates"| node_flows_classic
node_tests_dist -->|"validates"| node_dist_base
node_tests_dist -->|"validates"| node_dist_cond

click node_core "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/core.py"
click node_dist_base "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/base.py"
click node_dist_cond "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/encoder.py"
click node_dist_target "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/target.py"
click node_dist_mh "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/mh_proposal.py"
click node_flows_base "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/base.py"
click node_flows_affine "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/affine/__init__.py"
click node_flows_spline "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/neural_spline/__init__.py"
click node_flows_classic "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/planar.py"
click node_flows_struct "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/mixing.py"
click node_nets "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/nets/mlp.py"
click node_utils "https://github.com/antsx/antsnormalizingflows/tree/main/src/antsnormflows/utils"
click node_transforms "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/transforms.py"
click node_sampling "https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/sampling/hais.py"
click node_examples "https://github.com/antsx/antsnormalizingflows/tree/main/examples"
click node_docs "https://github.com/antsx/antsnormalizingflows/blob/main/docs/index.md"
click node_tests_core "https://github.com/antsx/antsnormalizingflows/blob/main/tests/core_test.py"
click node_tests_flows "https://github.com/antsx/antsnormalizingflows/tree/main/tests/flows"
click node_tests_dist "https://github.com/antsx/antsnormalizingflows/tree/main/tests/distributions"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_core,node_dist_base,node_dist_cond,node_dist_target,node_dist_mh,node_flows_base,node_flows_affine,node_flows_spline,node_flows_classic,node_flows_struct,node_sampling toneBlue
class node_nets,node_utils,node_transforms toneAmber
class node_examples,node_docs toneMint
class node_tests_core,node_tests_flows,node_tests_dist toneRose
```

## Components

### Library code

Primary PyTorch package under src/antsnormflows.

- [NormalizingFlow](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/core.py): model orchestrator · Central flow model that composes a base distribution with ordered transforms and handles forward/inverse density paths. · `src/antsnormflows/core.py`
- [Base dist](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/base.py): distribution core · Base and prior distributions that seed latent-space sampling and density evaluation. · `src/antsnormflows/distributions/base.py`
- [Encoder/decoder](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/encoder.py): learned distributions · Latent-variable distribution components used for conditional inference and generative modeling. · `src/antsnormflows/distributions/encoder.py`
- [Target dist](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/target.py): target density · Target-distribution helpers for approximation and probabilistic modeling workflows. · `src/antsnormflows/distributions/target.py`
- [Proposal tools](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/distributions/mh_proposal.py): mh proposal · Metropolis-Hastings style proposal machinery for advanced sampling flows. · `src/antsnormflows/distributions/mh_proposal.py`
- [Flow base](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/base.py): flow interface · Shared flow abstractions and base contracts for invertible tensor transforms. · `src/antsnormflows/flows/base.py`
- [Affine flows](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/affine/__init__.py): flow family · Coupling, autoregressive, and Glow-style affine transforms. · `src/antsnormflows/flows/affine/__init__.py`
- [Spline flows](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/neural_spline/__init__.py): flow family · Neural spline coupling and autoregressive transforms plus wrappers. · `src/antsnormflows/flows/neural_spline/__init__.py`
- [Classical flows](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/planar.py): flow family · Planar, radial, residual, periodic, and stochastic flow transforms. · `src/antsnormflows/flows/planar.py`
- [Structure ops](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/flows/mixing.py): tensor transforms · Mixing, normalization, and reshape transforms that organize flow topology and tensor shape. · `src/antsnormflows/flows/mixing.py`
- [HAIS sampling](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/sampling/hais.py): auxiliary sampler · 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](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/nets/mlp.py): neural builders · Reusable MLP, CNN, MADE, ResNet, and Lipschitz networks that parameterize flow and distribution blocks. · `src/antsnormflows/nets/mlp.py`
- [Utilities](https://github.com/antsx/antsnormalizingflows/tree/main/src/antsnormflows/utils): shared helpers · Masks, spline math, preprocessing, evaluation, neural-network helpers, and optimization support. · `src/antsnormflows/utils`
- [Transforms API](https://github.com/antsx/antsnormalizingflows/blob/main/src/antsnormflows/transforms.py): shared transform utils · General transform utilities shared across subpackages. · `src/antsnormflows/transforms.py`

### Examples and docs

Notebooks, scripts, and documentation showing intended composition and use.

- [Examples](https://github.com/antsx/antsnormalizingflows/tree/main/examples): usage demos · Notebooks and scripts demonstrating base distributions, flow stacks, conditional models, and VAE-style usage. · `examples`
- [Docs](https://github.com/antsx/antsnormalizingflows/blob/main/docs/index.md): documentation · MkDocs-based project documentation and generated reference material. · `docs/index.md`

### Tests and CI

Mirrored tests and automation for validation and release.

- [Core tests](https://github.com/antsx/antsnormalizingflows/blob/main/tests/core_test.py): test suite · Tests for model composition and core behavior. · `tests/core_test.py`
- [Flow tests](https://github.com/antsx/antsnormalizingflows/tree/main/tests/flows): test suite · Mirrored tests for affine, spline, classical, and structural flow families. · `tests/flows`
- [Dist tests](https://github.com/antsx/antsnormalizingflows/tree/main/tests/distributions): test suite · Mirrored tests for base, encoder, decoder, prior, and target distribution logic. · `tests/distributions`

## Connections

- 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.