jaxknn 0.1.0__tar.gz

jaxknn-0.1.0/.github/workflows/build_docs.yml

@@ -0,0 +1,86 @@
+name: Build docs
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+  pull_request:
+
+concurrency:
+  # Each push to main gets its own group, so merges never cancel each other
+  # and main keeps a complete CI record. All other events keep the per-ref
+  # group and still supersede older runs.
+  group: >-
+    ${{ github.workflow }}-${{ github.event_name == 'push' && github.sha || github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+env:
+  FORCE_COLOR: 3
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
+        with:
+          python-version: "3.14"
+
+      - name: Install dependencies
+        run: uv sync --no-default-groups --group docs --locked
+
+
+      - name: Build docs
+        run: uv run --frozen zensical build --clean
+
+  upload:
+    needs: build
+    if: github.ref == 'refs/heads/main'
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      actions: write
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
+        with:
+          python-version: "3.14"
+
+      - name: Install dependencies
+        run: uv sync --no-default-groups --group docs --locked
+
+
+      - name: Build docs
+        run: uv run --frozen zensical build --clean
+
+      - name: Upload pages artifact
+        uses: actions/upload-pages-artifact@fc324d3547104276b827a68afc52ff2a11cc49c9 # v5.0.0
+        with:
+          path: site
+
+  deploy:
+    needs: upload
+    if: github.ref == 'refs/heads/main'
+    runs-on: ubuntu-latest
+    permissions:
+      pages: write
+      id-token: write
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@368f82528645a54fb793d4d04e342629a3f51346 # v5.0.1

jaxknn-0.1.0/.github/workflows/cd.yml

@@ -0,0 +1,53 @@
+name: CD
+
+on:
+  workflow_dispatch:
+  push:
+    tags:
+      - "v*"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: false
+
+env:
+  FORCE_COLOR: 3
+
+jobs:
+  build:
+    name: Build & inspect package
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      id-token: write
+      attestations: write
+
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+          persist-credentials: false
+
+      - uses: hynek/build-and-inspect-python-package@2abe76da66d0a6a4a227101f9348ee855797cfa5 # v3.0.1
+        with:
+          attest-build-provenance-github: true
+
+  publish-pypi:
+    name: Publish to PyPI
+    needs: [build]
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/jaxknn
+    permissions:
+      id-token: write
+    if: startsWith(github.ref, 'refs/tags/v')
+
+    steps:
+      - name: Download built artifact to dist/
+        uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
+        with:
+          name: Packages
+          path: dist
+
+      - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2

jaxknn-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,145 @@
+name: CI
+
+on:
+  workflow_dispatch:
+  pull_request:
+  push:
+    branches:
+      - main
+  schedule:
+    - cron: "0 0 * * 4" # Every Thursday at midnight UTC
+
+permissions:
+  contents: read
+
+concurrency:
+  # Each push to main gets its own group, so merges never cancel each other
+  # and main keeps a complete CI record. All other events keep the per-ref
+  # group and still supersede older runs.
+  group: >-
+    ${{ github.workflow }}-${{ github.event_name == 'push' && github.sha || github.ref }}
+  cancel-in-progress: true
+
+env:
+  FORCE_COLOR: 3
+
+jobs:
+  lint:
+    name: Lint
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
+
+      - name: Install
+        run: uv sync --no-default-groups --group lint --locked
+
+      - name: Checks with prek
+        # Not a real commit -- no-commit-to-branch would always fail here.
+        run: uv run --frozen prek run --all-files
+        env:
+          SKIP: no-commit-to-branch
+
+  tests:
+    name: Python ${{ matrix.python-version }} on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    needs: [lint]
+    permissions:
+      contents: read
+      id-token: write # Codecov OIDC upload
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12", "3.14"]
+        os: [ubuntu-latest, macos-latest, windows-latest]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install
+        run: uv sync --no-default-groups --group tests --locked
+
+      - name: Test with pytest
+        run: uv run --frozen pytest --cov --cov-report=xml --cov-report=term
+
+      - name: Upload coverage report
+        uses: codecov/codecov-action@303a32d7a59b442fa8d48b6a1cc6825c09c847a5 # v7.1.1
+        with:
+          use_oidc: true
+
+  oldest-supported:
+    name: Oldest supported deps
+    runs-on: ubuntu-latest
+    needs: [lint]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
+        with:
+          python-version: "3.12"
+
+      - name: Install the project
+        run: uv sync --no-default-groups --group tests --resolution lowest-direct
+
+      - name: Test package
+        run: uv run --no-sync pytest
+
+  newest-supported:
+    name: Newest supported deps
+    runs-on: ubuntu-latest
+    needs: [lint]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
+        with:
+          python-version: "3.14"
+
+      - name: Install the project
+        run: uv sync --no-default-groups --group tests --upgrade
+
+      - name: Test package
+        run: uv run --no-sync pytest
+
+  benchmark-smoke:
+    name: Benchmark smoke test
+    runs-on: ubuntu-latest
+    needs: [lint]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
+
+      - name: Install
+        run: uv sync --no-default-groups --group bench --locked
+
+      - name: Run the benchmarks once
+        run: uv run --frozen pytest tests/benchmarks --override-ini="addopts=" -p no:randomly
+
+  ci-pass:
+    name: CI Pass
+    runs-on: ubuntu-latest
+    needs: [lint, tests, oldest-supported, newest-supported, benchmark-smoke]
+    if: always()
+    steps:
+      - name: All required jobs passed
+        uses: re-actors/alls-green@b5b5b37504aa4183270bd3d855c52a67f212be35 # v1.3.0
+        with:
+          jobs: ${{ toJSON(needs) }}

jaxknn-0.1.0/.gitignore

@@ -0,0 +1,14 @@
+__pycache__/
+*.py[cod]
+.venv/
+build/
+dist/
+*.egg-info/
+.coverage
+coverage.xml
+.pytest_cache/
+.ruff_cache/
+.hypothesis/
+.codspeed/
+site/
+src/jaxknn/_version.py

jaxknn-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,42 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: "v6.0.0"
+    hooks:
+      - id: no-commit-to-branch
+        args: ["--branch", "main"]
+        always_run: true
+      - id: check-added-large-files
+      - id: check-case-conflict
+      - id: check-merge-conflict
+      - id: check-yaml
+      - id: debug-statements
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+      - id: name-tests-test
+        args: ["--pytest-test-first"]
+        exclude: ^tests/(benchmarks/|_compat\.py$)
+      - id: trailing-whitespace
+
+  - repo: https://github.com/python-jsonschema/check-jsonschema
+    rev: 0.36.0
+    hooks:
+      - id: check-github-workflows
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.14
+    hooks:
+      - id: ruff-check
+        types_or: [python, pyi]
+        args: [--fix, --show-fixes]
+      - id: ruff-format
+        types_or: [python, pyi]
+
+  - repo: https://github.com/codespell-project/codespell
+    rev: "v2.4.1"
+    hooks:
+      - id: codespell
+
+  - repo: https://github.com/ComPWA/taplo-pre-commit
+    rev: v0.9.3
+    hooks:
+      - id: taplo-format

jaxknn-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

jaxknn-0.1.0/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2025-2026, Nathaniel Starkman.
+All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jaxknn-0.1.0/PKG-INFO

@@ -0,0 +1,97 @@
+Metadata-Version: 2.5
+Name: jaxknn
+Version: 0.1.0
+Summary: Exact k-nearest neighbours in JAX: a kd-tree that traces under jit, vmap and grad.
+Project-URL: Bug Tracker, https://github.com/nstarman/jaxknn/issues
+Project-URL: Changelog, https://github.com/nstarman/jaxknn/releases
+Project-URL: Documentation, https://nstarkman.space/jaxknn
+Project-URL: Source Code, https://github.com/nstarman/jaxknn
+Author-email: Nathaniel Starkman <nstarman@users.noreply.github.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: jax,kd-tree,knn,nearest-neighbors
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering
+Classifier: Typing :: Typed
+Requires-Python: >=3.12
+Requires-Dist: equinox>=0.13.5
+Requires-Dist: jax>=0.7.2
+Requires-Dist: jaxtyping>=0.3.5
+Requires-Dist: numpy>=2.0
+Description-Content-Type: text/markdown
+
+# jaxknn
+
+Exact k-nearest neighbours in [JAX](https://github.com/jax-ml/jax): a kd-tree that traces under `jit`, `vmap` and `grad`.
+
+```bash
+pip install jaxknn
+```
+
+```python
+>>> import jax
+>>> import jax.numpy as jnp
+>>> import jaxknn
+
+>>> points = jax.random.normal(jax.random.key(0), (1000, 3))
+>>> idx, dist = jaxknn.knn(points, 5)  # each point's 5 nearest *other* points
+>>> idx.shape, dist.shape
+((1000, 5), (1000, 5))
+
+>>> queries = jnp.zeros((2, 3))
+>>> idx, dist = jax.jit(lambda p, q: jaxknn.knn(p, 1, queries=q))(points, queries)
+>>> idx.shape
+(2, 1)
+
+```
+
+## What it gives you
+
+- **Exact results.** No approximation. Ties go to the lower index, identically eagerly and under `jit`.
+- **Fully traceable.** `knn` works inside `jit`, `vmap` and `grad`. Distances are differentiable with respect to the coordinates, with the neighbour selection held fixed, and the gradient is 0 rather than NaN at coincident points.
+- **Robust inputs.**
+  - NaN or inf input is an error, not a silently wrong answer: `ValueError` eagerly, a runtime error under `jit`.
+  - Integers and bfloat16 are promoted to float32.
+  - Coordinates are divided by an exact power of two when float32 would overflow or underflow, so values from about 1e-38 to 3e38 keep their neighbours. Typical cases are metres at kiloparsec scale, or tiny lengths.
+- **Compiles once per size bucket.** Eager calls pad `n` to a bucket (2ʲ or 1.5·2ʲ), so inputs of different sizes share compiled code.
+- **A raw tree for custom queries.** `jaxknn.tree` exposes `build_tree`, `knn`, `all_knn` and label-excluding queries (`knn(tree, q, k, exclude=(point_labels, query_labels))`: the nearest point with a *different* label, with whole same-label subtrees pruned).
+
+`jaxknn.brute_knn` has the same contract by brute force. It's faster for small inputs.
+
+## When to use it
+
+Use jaxknn when the kNN has to live inside a traced JAX program: under `jit`, `vmap` or `grad`, or on an accelerator.
+
+On a CPU, eagerly, SciPy's `cKDTree` is faster, because it queries with every core. Here's all-points kNN (k=10, a 3-D stream with 2% interlopers), comparing jaxknn compiled under `jit` with SciPy called eagerly. It was measured on an Apple M2 Pro under load, so the ratios matter more than the absolute times.
+
+| n | `jit` compile (first call) | jaxknn under `jit`, steady | SciPy, eager | jaxknn / SciPy |
+|---|---|---|---|---|
+| 2,000 | 9.4 s | 17.8 ms | 6.0 ms | 2.9× |
+| 10,000 | 11.6 s | 106 ms | 16.4 ms | 6.5× |
+| 100,000 | 13.2 s | 1,261 ms | 148 ms | 8.5× |
+| 1,000,000 | 26.8 s | 11,719 ms | 1,990 ms | 5.9× |
+
+Against [jaxkd](https://github.com/dodgebc/jaxkd) at n = 100,000 (all-points, k = 10, CPU):
+
+| data | jaxknn | jaxkd | SciPy |
+|---|---|---|---|
+| stream | 307 ms | 815 ms | 73 ms |
+| blob | 355 ms | 1,161 ms | 80 ms |
+| stream + interlopers | 513 ms | ~140 s | 78 ms |
+
+## Origin
+
+jaxknn was extracted from [phasecurvefit](https://github.com/GalacticDynamics/phasecurvefit), where it is the traceable kNN behind `MSTOrderer`, with its history. See [phasecurvefit#197](https://github.com/GalacticDynamics/phasecurvefit/issues/197).
+
+## License
+
+MIT

jaxknn-0.1.0/README.md

@@ -0,0 +1,66 @@
+# jaxknn
+
+Exact k-nearest neighbours in [JAX](https://github.com/jax-ml/jax): a kd-tree that traces under `jit`, `vmap` and `grad`.
+
+```bash
+pip install jaxknn
+```
+
+```python
+>>> import jax
+>>> import jax.numpy as jnp
+>>> import jaxknn
+
+>>> points = jax.random.normal(jax.random.key(0), (1000, 3))
+>>> idx, dist = jaxknn.knn(points, 5)  # each point's 5 nearest *other* points
+>>> idx.shape, dist.shape
+((1000, 5), (1000, 5))
+
+>>> queries = jnp.zeros((2, 3))
+>>> idx, dist = jax.jit(lambda p, q: jaxknn.knn(p, 1, queries=q))(points, queries)
+>>> idx.shape
+(2, 1)
+
+```
+
+## What it gives you
+
+- **Exact results.** No approximation. Ties go to the lower index, identically eagerly and under `jit`.
+- **Fully traceable.** `knn` works inside `jit`, `vmap` and `grad`. Distances are differentiable with respect to the coordinates, with the neighbour selection held fixed, and the gradient is 0 rather than NaN at coincident points.
+- **Robust inputs.**
+  - NaN or inf input is an error, not a silently wrong answer: `ValueError` eagerly, a runtime error under `jit`.
+  - Integers and bfloat16 are promoted to float32.
+  - Coordinates are divided by an exact power of two when float32 would overflow or underflow, so values from about 1e-38 to 3e38 keep their neighbours. Typical cases are metres at kiloparsec scale, or tiny lengths.
+- **Compiles once per size bucket.** Eager calls pad `n` to a bucket (2ʲ or 1.5·2ʲ), so inputs of different sizes share compiled code.
+- **A raw tree for custom queries.** `jaxknn.tree` exposes `build_tree`, `knn`, `all_knn` and label-excluding queries (`knn(tree, q, k, exclude=(point_labels, query_labels))`: the nearest point with a *different* label, with whole same-label subtrees pruned).
+
+`jaxknn.brute_knn` has the same contract by brute force. It's faster for small inputs.
+
+## When to use it
+
+Use jaxknn when the kNN has to live inside a traced JAX program: under `jit`, `vmap` or `grad`, or on an accelerator.
+
+On a CPU, eagerly, SciPy's `cKDTree` is faster, because it queries with every core. Here's all-points kNN (k=10, a 3-D stream with 2% interlopers), comparing jaxknn compiled under `jit` with SciPy called eagerly. It was measured on an Apple M2 Pro under load, so the ratios matter more than the absolute times.
+
+| n | `jit` compile (first call) | jaxknn under `jit`, steady | SciPy, eager | jaxknn / SciPy |
+|---|---|---|---|---|
+| 2,000 | 9.4 s | 17.8 ms | 6.0 ms | 2.9× |
+| 10,000 | 11.6 s | 106 ms | 16.4 ms | 6.5× |
+| 100,000 | 13.2 s | 1,261 ms | 148 ms | 8.5× |
+| 1,000,000 | 26.8 s | 11,719 ms | 1,990 ms | 5.9× |
+
+Against [jaxkd](https://github.com/dodgebc/jaxkd) at n = 100,000 (all-points, k = 10, CPU):
+
+| data | jaxknn | jaxkd | SciPy |
+|---|---|---|---|
+| stream | 307 ms | 815 ms | 73 ms |
+| blob | 355 ms | 1,161 ms | 80 ms |
+| stream + interlopers | 513 ms | ~140 s | 78 ms |
+
+## Origin
+
+jaxknn was extracted from [phasecurvefit](https://github.com/GalacticDynamics/phasecurvefit), where it is the traceable kNN behind `MSTOrderer`, with its history. See [phasecurvefit#197](https://github.com/GalacticDynamics/phasecurvefit/issues/197).
+
+## License
+
+MIT

jaxknn-0.1.0/conftest.py

@@ -0,0 +1,19 @@
+"""Doctest the code blocks in the README, the docs and the docstrings."""
+
+from doctest import ELLIPSIS, NORMALIZE_WHITESPACE
+
+from sybil import Sybil
+from sybil.parsers import markdown, rest
+
+_flags = ELLIPSIS | NORMALIZE_WHITESPACE
+
+pytest_collect_file = (
+    Sybil(
+        parsers=[
+            markdown.PythonCodeBlockParser(doctest_optionflags=_flags),
+            markdown.SkipParser(),
+        ],
+        patterns=["*.md"],
+    )
+    + Sybil(parsers=[rest.DocTestParser(optionflags=_flags)], patterns=["*.py"])
+).pytest()

jaxknn-0.1.0/docs/_overrides/partials/source.html

@@ -0,0 +1,10 @@
+{% import "partials/language.html" as lang with context %}
+<a href="{{ config.repo_url }}" title="{{ lang.t('source.link.title') }}" class="md-source" data-md-component="source">
+  <div class="md-source__icon md-icon">
+    {% set icon = config.theme.icon.repo or "fontawesome/brands/git-alt" %}
+    {% include ".icons/" ~ icon ~ ".svg" %}
+  </div>
+  <div class="md-source__repository">
+    {{ config.repo_name }}
+  </div>
+</a>