pyproject.toml

[tool.poetry]
name = "approx-topk"
version = "0.1.0"
description = "Bucketed top-k for PyTorch using a priority queue"
authors = ["Oscar Key <oscar.t.key@gmail.com>"]
readme = "README.md"
license = "MIT"
include = ["approx_topk/*.cu"]

[tool.poetry.dependencies]
python = "^3.10"
torch = { version = "^2.4.0" }
setuptools = "^72.1.0"

[tool.poetry.group.dev.dependencies]
mypy = "^1.11.1"
pytest = "^8.3.2"
types-tqdm = "^4.66.0.20240417"
ruff = "^0.8.0"

[tool.poetry.group.benchmarks.dependencies]
matplotlib = "^3.9.0"
tqdm = "^4.66.5"
ipykernel = "^6.29.5"
pandas = "^2.2.2"
seaborn = "^0.13.2"
pylibraft-cu12 = { version = "^24.10.0", source = "nvidia" }
cupy-cuda12x = "^13.3.0"

[[tool.poetry.source]]
name = "pytorch_cu124"
url = "https://download.pytorch.org/whl/cu124"
priority = "explicit"

[[tool.poetry.source]]
name = "nvidia"
url = "https://pypi.nvidia.com"
priority = "explicit"

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"