Installation

Install from PyPI

Install the core kpnn2 package from PyPI with:

pip install kpnn2

The core installation supports parsing an edgelist into a LayeredSpec, building MaskedLinear layers, aligning named inputs, and mapping a layer tensor back to named nodes. Training stays ordinary PyTorch. Supported architectures lists which architecture families you can assemble from those primitives.

kpnn2 requires Python 3.10 or later. See the tested Python versions in the badge on the Home page.

Core dependencies are torch, pandas, numpy, and xarray.

Attribution tools

kpnn2 does not run attribution for you. Captum is not a package dependency.

If you want Captum (as in Feedforward example, the Time-series example, and the Transformer example), install it yourself:

pip install captum

Then pass the resulting tensor to map_node_attributions. Fold named scores with aggregate_node_attributions when you want one value per node. Any other attribution method that yields a tensor of named-layer units works the same way. Extra axes (class, step, …) stay as extra xarray dimensions.

Development installation

To work on the package locally, clone the repository and install it in editable mode from the project root:

git clone git@github.com:Thomas-Rauter/kpnn2.git
cd kpnn2
pip install -e .

Optional dependency groups

Install development dependencies with:

pip install -e ".[dev]"

That extra pins the Ruff version CI uses. Run Ruff as python -m ruff, so a different ruff on PATH cannot stand in for the pinned one. The full list of local checks is dev/check_code_locally.txt. To upgrade Ruff, bump the pin in pyproject.toml, run python -m ruff format ., and commit the pin and any rewrites together.

Install documentation dependencies with:

pip install -e ".[docs]"

The docs extra includes the tools used by the documentation notebooks, such as Captum, Graphviz (Python package), seaborn, and Jupyter.

Notebooks and documentation

Some documentation notebooks draw graphs with Graphviz. If a notebook requires Graphviz rendering, you also need the system-level Graphviz installation, not only the Python package.

For example, on Ubuntu or Debian:

sudo apt install graphviz

Verify the installation

A minimal core-installation smoke test is:

python -c "import kpnn2; print('kpnn2 imported successfully')"