Skip to content

17. Contributing to NNx

Thanks for being interested in contributing. NNx is a small library; the goal is to keep it small, tested, and useful for the existing notebook consumers while inviting new ones.

1. Getting set up

git clone https://github.com/thekaveh/NNx.git
cd NNx
python -m pip install -r requirements-tools.txt
uv sync --all-extras --frozen
uv run pre-commit install       # optional but recommended

Diagram fallback generation also needs the native Cairo library: install it with brew install cairo on macOS or sudo apt-get install libcairo2 on Debian/Ubuntu. CairoSVG itself is pinned in the docs-publish dependency group.

Verify a clean baseline:

uv run pytest                          # full suite
uv run ruff check src/ tests/ examples/ scripts/  # lint
uv run ruff format --check src/ tests/ examples/ scripts/  # format check
uv run pyright --warnings              # type check
uv run python -m scripts.docs.extract_architecture_svg --check
uv run python -m scripts.docs.build_docs --check
uv run python -m scripts.docs.build_docs
uv run mkdocs build --strict           # generated docs (gates CI)

Repository Markdown is canonical. docs/manifest.yaml is the single page inventory and each source H1 begins with its manifest number; scripts.docs.build_docs generates ignored mkdocs.yml, generated/site, and generated/wiki projections. Do not edit generated outputs. A successful push to main publishes the same canonical content to GitHub Pages and the repository wiki.

NNx documentation projection

Diagram HTML masters first flow through scripts.docs.extract_architecture_svg. The site selects SVG assets while the repository and wiki select PNG fallbacks. CI checks canonical GitHub links, tracked-file ownership, placeholder markers, stale API signatures and diagrams, projection determinism, generated links, and a strict MkDocs build before publication.

Useful env vars:

  • NNX_TQDM_DISABLE=1 silences the training progress bar. Set this in CI / non-TTY contexts, and in any test that drives NNModel.train() or Trainer.train() (the test suite's conftest.py already does this session-wide). Accepts 1 / true / yes, case-insensitive.

2. Workflow

  1. Open an issue first for non-trivial changes — saves churn if the design is off. Tiny fixes can go straight to PR.
  2. Branch from develop. Name branches descriptively (fix/..., feat/..., docs/..., refactor/...).
  3. Write tests. Every PR that changes behavior should land with a focused test that fails on main and passes on the branch. The existing tests/test_*_series.py files (organized by audit pass) are good models.
  4. Keep PRs small. One coherent change per PR is much easier to review than a sweeping mix.

3. What we care about

  • Strict back-compat for the existing notebook consumer. Don't rename, remove, or restructure public APIs without a migration path. Preserve existing runs/<id>/ artifacts; on-disk evolution requires a versioned, backward-compatible reader. New fields on params dataclasses must omit themselves from .state() when set to their defaults (preserves run.id hashes). See the omit-when-default regression tests in tests/test_params_round_trip.py (search for test_nn_*_state_omits_*_when_*) for the canonical pattern.
  • State / from_state round-trip. Every params dataclass with a state() method must round-trip cleanly through from_state(state()). The contract is enforced by tests/test_params_round_trip.py.
  • Tests run on CPU and finish fast. Keep new tests under a few seconds; use small TensorDataset fixtures from tests/conftest.py.
  • One-line update to CHANGELOG.md under [Unreleased] for any user-visible change.

4. Style

  • Ruff enforces formatting and a curated lint rule set (E F W B I UP). Run ruff check --fix src/ tests/ examples/ scripts/ and ruff format src/ tests/ examples/ scripts/ before pushing. Pre-commit handles both automatically when installed.
  • Type annotations are encouraged on new code. We type-check with pyright (basic mode) in CI, with --strict planned over time.
  • Docstrings on public functions / classes explain the why (constraints, edge cases) — not just the what. Multi-paragraph is fine when warranted.
  • Comments explain non-obvious decisions, hidden constraints, or surprising behavior. Don't narrate code that's already self-documenting.

5. Testing

uv run pytest                          # full suite
uv run pytest tests/test_pass2_n_series.py::test_n7_evaluate_aggregates_across_batches
uv run pytest -k "graph"               # name filter
uv run pytest --cov=nnx --cov-report=term-missing  # with coverage

Tests live under tests/. The conftest.py registers a handful of hygiene fixtures (session-wide NNX_TQDM_DISABLE, a per-test env_snapshot cache reset, and a dynamo-dispatch skip guard); otherwise it's intentionally minimal. Add shared fixtures there when boilerplate repeats across multiple tests, not preemptively.

6. Submitting a PR

  • Push to your fork and open a PR against develop.
  • Fill in the PR template (Summary / Test plan).
  • Wait for CI to go green (lint + format + tests + mkdocs on Python 3.10 through 3.14).
  • Address review comments by pushing new commits — we squash on merge.

7. Releases

NNx uses release-please for automated version bumps, changelog updates, and tagging. Contributors don't touch versions or tags — just write a Conventional Commit-style PR title (feat:, fix:, chore:, docs:, etc.) and add a one-line entry under [Unreleased] in CHANGELOG.md for any user-visible change.

The end-to-end flow:

  1. Every merge to main updates a long-lived "Release" PR maintained by release-please.yml. The PR accumulates the next version + CHANGELOG.md diff based on the conventional-commit types since the last tag. Pre-1.0, feat: triggers a minor bump (0.X.0); fix: and most other types trigger a patch bump (0.X.Y).
  2. A maintainer reviews and merges the Release PR when ready to ship. Release Please explicitly dispatches the required CI and security checks for its managed branch, then creates the tag and draft release and dispatches the top-level release.yml workflow with the immutable release commit and tag.
  3. release.yml verifies the Release Please commit belongs to main, runs the full test matrix, builds the package, publishes through PyPI trusted publishing (gated by the pypi GitHub Environment's approval rule), and verifies pip install thekaveh-nnx==X.Y.Z from a clean environment. Keeping trusted publishing in this top-level workflow ensures its OIDC token and package attestations carry the same PyPI-authorized workflow identity.

The [project] version is intentionally static and managed by release-please. Do not distribute wheels or sdists built from an untagged commit: after development resumes, such an artifact can contain code newer than the release while still carrying the last release number. Use editable installs for local source work. Distributable artifacts must come from Release Please's dispatched release workflow, which verifies tag/version agreement before publishing; direct tag pushes do not publish packages.

8. Things we won't merge

  • Changes that break on-disk format compatibility without a versioned reader.
  • Public API renames without a deprecation shim and a __getattr__ alias for at least one minor version.
  • Code without tests.
  • Dependencies added to [project.dependencies] (the core deps list) when they could go under [project.optional-dependencies] instead.

9. License

By contributing you agree that your contribution will be licensed under the Apache License 2.0.