File size: 4,645 Bytes
ecc81b3 ede9e39 ecc81b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | [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"]
|