| [build-system] | |
| requires = ["hatchling"] | |
| build-backend = "hatchling.build" | |
| [project] | |
| name = "torch-dimensions" | |
| version = "0.3.1" | |
| description = "N-dimensional models for PyTorch: SSMs, RNNs, transformers and convolutions over arbitrary lattices" | |
| readme = "README.md" | |
| requires-python = ">=3.10" | |
| license = "Apache-2.0" | |
| license-files = ["LICENSE", "NOTICE"] | |
| authors = [{ name = "triadastra" }] | |
| keywords = ["pytorch", "state-space-models", "mamba", "s4", "axial-attention", "vision-transformer", "multidimensional"] | |
| classifiers = [ | |
| "Development Status :: 2 - Pre-Alpha", | |
| "Intended Audience :: Science/Research", | |
| "Programming Language :: Python :: 3.10", | |
| "Programming Language :: Python :: 3.11", | |
| "Programming Language :: Python :: 3.12", | |
| "Topic :: Scientific/Engineering :: Artificial Intelligence", | |
| ] | |
| # torch and nothing else at install time. The default S4/S4D/Mamba mixers run | |
| # the vendored upstream code, whose own dependencies ([upstream] below) are | |
| # auto-installed on the first call that needs them; portable=True models and | |
| # everything else in the library never need more than torch. | |
| dependencies = ["torch>=2.4"] | |
| [project.optional-dependencies] | |
| mamba = ["mamba-ssm>=2.0"] | |
| fla = ["flash-linear-attention"] | |
| # The vendored reference implementations (torch_dimensions._vendor) run the | |
| # original authors' code verbatim; these are the dependencies *their* files | |
| # import (einops everywhere; numpy+scipy for HiPPO; hydra+omegaconf because | |
| # upstream's S4Block builds its inner layer through their hydra registry). | |
| upstream = ["einops", "numpy", "scipy", "hydra-core", "omegaconf"] | |
| # A torch pickle executes arbitrary code when it is opened. `.safetensors` | |
| # checkpoints do not, and `td.save`/`td.load` pick the container from the | |
| # file extension. | |
| # safetensors' torch bindings import numpy at runtime, and torch itself does not | |
| # require it — so the extra says so rather than failing three frames deep. | |
| safetensors = ["safetensors>=0.4", "numpy"] | |
| dev = [ | |
| "pytest>=8", "pytest-cov", "ruff>=0.6", "mypy", "hypothesis", "pyyaml", | |
| # Optional at runtime, mandatory for testing the paths that use them — | |
| # otherwise CI skips them and reports green (DEBUG.md #24). The [upstream] | |
| # deps are here so tests/test_vendored.py actually runs the vendored | |
| # originals. | |
| "safetensors>=0.4", "numpy", "einops", "scipy", "hydra-core", "omegaconf", | |
| # benchmarks/figure.py draws the published device-comparison sheet. Here | |
| # rather than optional so the figure is regenerated and checked, not | |
| # hand-maintained alongside numbers it is supposed to be derived from. | |
| "matplotlib>=3.8", | |
| ] | |
| all = ["torch-dimensions[mamba,fla,safetensors,upstream]"] | |
| [project.urls] | |
| Homepage = "https://github.com/triadastra/torch-dimensions" | |
| Source = "https://github.com/triadastra/torch-dimensions" | |
| [tool.hatch.build] | |
| # The viewer bundle is a build artifact: gitignored, rebuilt by CI before | |
| # packaging, and listed here because hatchling otherwise honours .gitignore | |
| # and would ship a wheel whose `td.viz.show` has nothing to serve. | |
| artifacts = ["src/torch_dimensions/viz/static/**"] | |
| [tool.hatch.build.targets.sdist] | |
| # The sdist is the library, not the repo: without this, viewer/node_modules | |
| # (35 MB of JS toolchain) rides along into every pip install from source. | |
| include = [ | |
| "src", | |
| "tests", | |
| "examples", | |
| "*.md", | |
| "LICENSE", | |
| ] | |
| [tool.hatch.build.targets.wheel] | |
| packages = ["src/torch_dimensions"] | |
| [tool.pytest.ini_options] | |
| testpaths = ["tests"] | |
| markers = [ | |
| "gpu: requires a CUDA device (deselect with '-m \"not gpu\"')", | |
| "mamba: requires the [mamba] extra", | |
| ] | |
| addopts = "-m 'not gpu'" | |
| [tool.coverage.run] | |
| source = ["torch_dimensions"] | |
| # Upstream's code is verified byte-for-byte against their repositories, not | |
| # line-covered by our tests — coverage of their internals is their business. | |
| omit = ["*/torch_dimensions/_vendor/*"] | |
| [tool.coverage.report] | |
| # Measured at 97% when the floor was set. The floor exists to catch a module | |
| # arriving with no tests at all, not to chase the last few lines — the | |
| # uncovered remainder is mostly `td.viz.show`'s blocking browser path, which a | |
| # test cannot exercise without opening a browser. | |
| fail_under = 95 | |
| show_missing = true | |
| [tool.ruff] | |
| line-length = 100 | |
| src = ["src", "tests"] | |
| # Vendored upstream files must stay byte-faithful to their sources; linting | |
| # them would demand edits the byte-diff test exists to forbid. | |
| extend-exclude = ["src/torch_dimensions/_vendor"] | |
| [tool.ruff.lint] | |
| select = ["E", "F", "W", "I", "UP", "B", "SIM", "NPY"] | |
| [tool.mypy] | |
| python_version = "3.10" | |
| files = ["src"] | |
| ignore_missing_imports = true | |
| exclude = ["src/torch_dimensions/_vendor"] | |