Metadata-Version: 2.4
Name: bonsai-nn-library
Version: 0.7.1
Summary: Common PyTorch utilities for research projects on neural networks.
Author-email: Richard Lange <opensource@bonsai-neuro-ai.com>
License: MIT
Project-URL: Repository, https://github.com/bonsai-neuro-ai/nn-library
Classifier: Programming Language :: Python :: 3.13
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Development Status :: 1 - Planning
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.13
Description-Content-Type: text/markdown
License-File: LICENSE.txt
Requires-Dist: jsonargparse
Requires-Dist: matplotlib
Requires-Dist: mlflow
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: parse
Requires-Dist: pycocotools
Requires-Dist: pydot
Requires-Dist: pylint>=4.0.7
Requires-Dist: pyyaml
Requires-Dist: scikit-learn
Requires-Dist: scipy
Requires-Dist: seaborn
Requires-Dist: tabulate
Requires-Dist: torchmetrics
Requires-Dist: tqdm
Provides-Extra: dev
Requires-Dist: coverage; extra == "dev"
Requires-Dist: coverage-badge; extra == "dev"
Requires-Dist: interrogate; extra == "dev"
Provides-Extra: cpu
Requires-Dist: torch>=2.11.0; extra == "cpu"
Requires-Dist: torchvision>=0.26.0; extra == "cpu"
Provides-Extra: cu130
Requires-Dist: torch>=2.11.0; extra == "cu130"
Requires-Dist: torchvision>=0.26.0; extra == "cu130"
Dynamic: license-file

# NN-Library

![Test coverage](badges/coverage.svg)
![Docstring coverage](badges/interrogate_badge.svg)

We in the [BONSAI Lab](https://bonsai-neuro-ai.com) do research on neural networks, among other
things, that requires loading/training/reconfiguring neural network models. This library is a
work-in-progress suite of in-house tools to address some pain-points we've encountered in our
research workflow.

We make no guarantees about the stability or usability of this library, but we hope that it can be
useful to others in the research community. If you have any questions or suggestions, please feel
free to reach out to us or open an issue on the GitHub repository.

## Installation

Using `pip`:

    pip install bonsai-nn-library[cpu]  # specifies torch-on-cpu
    pip install bonsai-nn-library[cu130]  # specifies torch-on-cuda-13.0

Using [`uv`](https://docs.astral.sh/uv/) (in a project):

    uv add bonsai-nn-library --extra cpu  # for cpu
    uv add bonsai-nn-library --extra cu130  # for cuda 13.0

__CUDA versioning note:__ configuring CUDA dependencies for pytorch is notoriously tricky. More
generally, `pip` and `uv` were not designed to handle switching between different dependencies on
host systems with different hardware or driver support. At time of writing this, there are some
open PyTorch issues and open PEPs that will someday make life easier. For now, our `pyproject.toml`
file supports `cpu` and `cu130`. For other If you have a different CUDA version, you could update
our `pyproject.toml`. See the disclaimer above: we don't guarantee our library will work for
everyone on all systems, but if you have a way to improve it we welcome contributions.

## Usage

The top-level import is `nn_lib`. Say you want to use one of our "fancy layers" like a low-rank
convolution. You can do so like this:

```python
from torch import nn
from nn_lib.models.fancy_layers import LowRankConv2d

model = nn.Sequential(
    LowRankConv2d(in_channels=3, out_channels=64, kernel_size=3, rank=8),
    nn.ReLU(),
    nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3),
    nn.ReLU(),
    nn.Flatten(),
    nn.LazyLinear(10)
)
```

## Useful thing \#1: improved `GraphModule`s.

PyTorch was not originally designed to handle explicit computation graphs, but it was added somewhat
later in the `torch.fx` module. Others might use tensorflow or jax for this, but we like PyTorch.
The `torch.fx.GraphModule` class is the built-in way to handle computation graphs in PyTorch, but it
lacks some features that we find useful. We have extended the `GraphModule` class in our
`GraphModulePlus` class, which inherits from `GraphModule` and adds some further functionality.

A motivating use-case is that we want to be able to "stitch" models together or extract out hidden
layer activity. This is a little tricky to get right using `GraphModule` alone, but we've added some
utilities like

* `GraphModulePlus.set_output(layer_name)`: use this to chop off the head of a model and make it
  output from a specific layer.
* `GraphModulePlus.new_from_merge(...)`: use this to merge or "stitch" existing models together. See
  `demos/demo_stitching.py` for a worked out example.

We've also done some metaprogramming trickery so that if you import `GraphModulePlus` anywhere in
your code, it will automatically inject itself into the `torch.fx` module. The surprising but
convenient behavior is:

```python
from torch import nn
from torch.fx import symbolic_trace
from nn_lib.models import GraphModulePlus

my_regular_torch_model = nn.Sequential(
    nn.Conv2d(3, 64, 3),
    nn.ReLU(),
    nn.Conv2d(64, 64, 3),
    nn.ReLU(),
    nn.Flatten(),
    nn.LazyLinear(10)
)

# Natively, symbolic_trace is expected to return a GraphModule, but we've injected GraphModulePlus
graphified_model = symbolic_trace(my_regular_torch_model)
assert isinstance(graphified_model, GraphModulePlus)
```

## Useful thing \#2: Fancy layers.

We have implemented a few "fancy" layers, available via `nn_lib.models` or
`nn_lib.models.fancy_layers` that we find useful in our research. These include:

* `Regressable` linear layers: a `Protocol` that allows linear layers to be initialized by least
  squares regression. This is useful for initializing a linear layer to approximate a function
  learned by a different model.
* `RegressableLinear`: a regressable version of `nn.Linear`
* `LowRankLinear`: a regressable linear layer with a low-rank factorization.
* `ProcrustesLinear`: a regressable linear layer constrained to rotation, with optional shift (bias)
  and optional scaling.
* A conv2d version of each of the above.

## Useful thing \#3: MLFLOW utilities.

We use MLFlow to track our experiments. We have a few utilities in `nn_lib.utils.mlfow_cli` that 
remove a bit of boilerplate from our code. The biggest contribution here is the `run_registry` which
makes it relatively easy to manage experiments where you want to submit a singleton mlflow run per 
unique set of parameters.

## Useful thing \#4: Linear algebra and regression helpers

See `nn_lib.utils.pca` for a variety of tools for analyzing linear subspaces such as effective
dimensionality and calculating principal components from data that may have missing values.

See `nn_lib.utils.stats` for helpers calculating variances and covariances, such as Welford's
algorithm for numerically-stable batch-wise streaming updates of means and variances.

See `nn_lib.utils.xval_nuc_norm` for some novel methods we're developing to calculate *cross-validated*
nuclear norms of cross-covariance matrices. It's useful for neural (dis)similarity analyses.

See `nn_lib.analysis.regression` for linear regression utilities such as regressing from `x` to `y`
from streamed/batched data. This is used extensively in `fancy_layers` where we support initializing
`Linear` or `Conv2d` layers by regressing to expected outputs. `demos/demo_stitching.py` shows off
this functionality.

## Useful thing \#5: NTK utilities.

See `nn_lib.analysis.ntk` for some neural tangent kernel utilities.

## Forthcoming/Planned features

* More fancy layers
* Vector Quantization utilities (but see `nn_lib.models.sparse_auto_encoder` which has some already)
* Further analysis utilities especially focused on calculating neural similarity measures.

## Obsolete/deprecated features

* lightning training and overly-complex CLI utilities. Some straggler files might still need to be
  cleaned up.

## Test and documentation coverage

We track test coverage (via `coverage.py`) and docstring coverage (via `interrogate`) for the
`badges/coverage.svg` and `badges/interrogate_badge.svg` badges at the top of this file. Since the
test suite requires CUDA, these aren't run in GitHub CI — instead, regenerate them locally (e.g. on
the lab server) after making changes and commit the updated SVGs:

    uv sync --extra dev
    scripts/update_coverage_badges.sh
