tenspec 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tenspec-0.1.0/.gitignore +22 -0
- tenspec-0.1.0/LICENSE +21 -0
- tenspec-0.1.0/PKG-INFO +116 -0
- tenspec-0.1.0/README.md +81 -0
- tenspec-0.1.0/examples/forecast_scale.py +92 -0
- tenspec-0.1.0/examples/locos_capture.py +161 -0
- tenspec-0.1.0/pyproject.toml +61 -0
- tenspec-0.1.0/src/tenspec/__init__.py +39 -0
- tenspec-0.1.0/src/tenspec/errors.py +89 -0
- tenspec-0.1.0/src/tenspec/numpy.py +343 -0
- tenspec-0.1.0/src/tenspec/py.typed +0 -0
- tenspec-0.1.0/src/tenspec/pydantic/__init__.py +4 -0
- tenspec-0.1.0/src/tenspec/pydantic/annotations.py +407 -0
- tenspec-0.1.0/src/tenspec/pydantic/calls.py +136 -0
- tenspec-0.1.0/src/tenspec/pydantic/models.py +82 -0
- tenspec-0.1.0/src/tenspec/pydantic/standalone.py +99 -0
- tenspec-0.1.0/src/tenspec/runtime/__init__.py +1 -0
- tenspec-0.1.0/src/tenspec/runtime/arrays.py +313 -0
- tenspec-0.1.0/src/tenspec/runtime/bindings.py +168 -0
- tenspec-0.1.0/src/tenspec/runtime/validation.py +368 -0
- tenspec-0.1.0/src/tenspec/torch.py +276 -0
- tenspec-0.1.0/src/tenspec/types/__init__.py +1 -0
- tenspec-0.1.0/src/tenspec/types/base.py +13 -0
- tenspec-0.1.0/src/tenspec/types/dtypes.py +140 -0
- tenspec-0.1.0/src/tenspec/types/operations.py +114 -0
- tenspec-0.1.0/src/tenspec/types/properties.py +74 -0
- tenspec-0.1.0/src/tenspec/types/shapes.py +289 -0
- tenspec-0.1.0/src/tenspec/types/tensor.py +28 -0
- tenspec-0.1.0/tests/integration/checked_caller.py +351 -0
- tenspec-0.1.0/tests/integration/pyrightconfig.json +6 -0
- tenspec-0.1.0/tests/integration/test_checkers.py +72 -0
- tenspec-0.1.0/tests/pydantic/test_annotations.py +294 -0
- tenspec-0.1.0/tests/pydantic/test_calls.py +229 -0
- tenspec-0.1.0/tests/pydantic/test_future_annotations.py +118 -0
- tenspec-0.1.0/tests/pydantic/test_models.py +211 -0
- tenspec-0.1.0/tests/pydantic/test_named_aliases.py +188 -0
- tenspec-0.1.0/tests/pydantic/test_standalone.py +183 -0
- tenspec-0.1.0/tests/runtime/test_arrays.py +143 -0
- tenspec-0.1.0/tests/runtime/test_bindings.py +117 -0
- tenspec-0.1.0/tests/runtime/test_validation.py +567 -0
- tenspec-0.1.0/tests/test_caller_example.py +366 -0
- tenspec-0.1.0/tests/test_imports.py +53 -0
- tenspec-0.1.0/tests/test_numpy.py +164 -0
- tenspec-0.1.0/tests/test_torch.py +170 -0
- tenspec-0.1.0/tests/types/test_dtypes.py +36 -0
- tenspec-0.1.0/tests/types/test_shapes.py +89 -0
- tenspec-0.1.0/tests/types/test_shared_grammar.py +72 -0
- tenspec-0.1.0/tests/types/test_tensor.py +39 -0
tenspec-0.1.0/.gitignore
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
*.egg-info/
|
|
6
|
+
build/
|
|
7
|
+
dist/
|
|
8
|
+
|
|
9
|
+
# Env + tool caches
|
|
10
|
+
.venv/
|
|
11
|
+
.env
|
|
12
|
+
.pytest_cache/
|
|
13
|
+
.ruff_cache/
|
|
14
|
+
.mypy_cache/
|
|
15
|
+
|
|
16
|
+
# Editor state
|
|
17
|
+
.vscode/
|
|
18
|
+
.idea/
|
|
19
|
+
.DS_Store
|
|
20
|
+
|
|
21
|
+
# Retained release evidence (visible, never committed)
|
|
22
|
+
artifacts/
|
tenspec-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Egor Dmitriev
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
tenspec-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: tenspec
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Tensor specifications that a function, a model, or a direct call checks at its boundary.
|
|
5
|
+
Project-URL: Source, https://github.com/egordm/tenspec
|
|
6
|
+
Project-URL: Documentation, https://egordm.github.io/tenspec/
|
|
7
|
+
Project-URL: Issues, https://github.com/egordm/tenspec/issues
|
|
8
|
+
Author: Egor Dmitriev
|
|
9
|
+
License-Expression: MIT
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
15
|
+
Classifier: Typing :: Typed
|
|
16
|
+
Requires-Python: >=3.13
|
|
17
|
+
Requires-Dist: pydantic<3,>=2.13.4
|
|
18
|
+
Requires-Dist: typing-extensions>=4.13
|
|
19
|
+
Provides-Extra: dev
|
|
20
|
+
Requires-Dist: numpy>=2.1; (python_version < '3.14') and extra == 'dev'
|
|
21
|
+
Requires-Dist: numpy>=2.3.2; (python_version >= '3.14') and extra == 'dev'
|
|
22
|
+
Requires-Dist: poethepoet>=0.30; extra == 'dev'
|
|
23
|
+
Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
24
|
+
Requires-Dist: ruff>=0.8; extra == 'dev'
|
|
25
|
+
Requires-Dist: torch>=2.6; (python_version < '3.14') and extra == 'dev'
|
|
26
|
+
Requires-Dist: torch>=2.9; (python_version >= '3.14') and extra == 'dev'
|
|
27
|
+
Requires-Dist: ty>=0.0.78; extra == 'dev'
|
|
28
|
+
Provides-Extra: numpy
|
|
29
|
+
Requires-Dist: numpy>=2.1; (python_version < '3.14') and extra == 'numpy'
|
|
30
|
+
Requires-Dist: numpy>=2.3.2; (python_version >= '3.14') and extra == 'numpy'
|
|
31
|
+
Provides-Extra: torch
|
|
32
|
+
Requires-Dist: torch>=2.6; (python_version < '3.14') and extra == 'torch'
|
|
33
|
+
Requires-Dist: torch>=2.9; (python_version >= '3.14') and extra == 'torch'
|
|
34
|
+
Description-Content-Type: text/markdown
|
|
35
|
+
|
|
36
|
+
# tenspec
|
|
37
|
+
|
|
38
|
+
Declare what an array must be, beside the value it describes, and check it at one boundary.
|
|
39
|
+
|
|
40
|
+
Tenspec moves the shape and dtype checks from the top of a function into its signature, where a
|
|
41
|
+
reader and a type checker both see them. It accepts or refuses the array the caller passed. It
|
|
42
|
+
converts nothing, moves nothing, and copies nothing, unless the declaration names a transform
|
|
43
|
+
that copies.
|
|
44
|
+
|
|
45
|
+
This is an early release, at version 0.1.0. The interface can change before version 1.0.
|
|
46
|
+
|
|
47
|
+
## Install
|
|
48
|
+
|
|
49
|
+
Install it with pip, or with uv. Either one is enough, so run the line for the tool you use:
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
pip install "tenspec[numpy]" # or: uv add "tenspec[numpy]"
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
Name the array library you need as an extra: `numpy` or `torch`. `import tenspec` loads neither
|
|
56
|
+
one. It needs Python 3.13 or newer.
|
|
57
|
+
|
|
58
|
+
## What it gives you
|
|
59
|
+
|
|
60
|
+
- **Ordinary types.** A declaration is an ordinary annotation carrying your backend's array
|
|
61
|
+
type, so a type checker reads it and untouched code still runs. The checkers, the settings
|
|
62
|
+
they need and their measured limits are on
|
|
63
|
+
[Type checkers and Ruff](https://egordm.github.io/tenspec/type-checkers.html).
|
|
64
|
+
- **Pydantic boundaries, natively.** `TensorContracts` relates the fields of your own model
|
|
65
|
+
inside one model validation.
|
|
66
|
+
- **Runtime checks are opt in.** An annotation alone performs no runtime work.
|
|
67
|
+
- **Transforms are explicit.** A value changes only when a declaration names a transform, and
|
|
68
|
+
the runtime verifies the result kept its class, shape, native dtype and device.
|
|
69
|
+
|
|
70
|
+
## The three entry points
|
|
71
|
+
|
|
72
|
+
| Entry point | Use it for |
|
|
73
|
+
|---|---|
|
|
74
|
+
| `@checked` | a function or a method, with its arguments and its return |
|
|
75
|
+
| `TensorContracts` | a Pydantic model, whose related fields share one boundary |
|
|
76
|
+
| `validate(value, annotation)` | one value, or one tuple of related values |
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
from typing import Literal as Shape
|
|
80
|
+
|
|
81
|
+
import numpy as np
|
|
82
|
+
|
|
83
|
+
from tenspec import checked
|
|
84
|
+
from tenspec.numpy import Float
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
@checked
|
|
88
|
+
def weigh_columns(
|
|
89
|
+
values: Float[Shape["rows features"]], weights: Float[Shape["features"]]
|
|
90
|
+
) -> Float[Shape["rows features"]]:
|
|
91
|
+
return values * weights
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
print(weigh_columns(np.ones((3, 4)), np.array([1.0, 2.0, 3.0, 4.0])).shape)
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
`weights` must cover as many features as `values` has columns. A disagreement is a run-time
|
|
98
|
+
refusal naming the axis, both extents, and the bindings the boundary held.
|
|
99
|
+
|
|
100
|
+
No type checker compares two shapes. A shape that disagrees is a run-time refusal, never a
|
|
101
|
+
static error.
|
|
102
|
+
|
|
103
|
+
## Documentation
|
|
104
|
+
|
|
105
|
+
[The documentation site](https://egordm.github.io/tenspec/) holds the guide, the API reference
|
|
106
|
+
and two executed tutorials. The
|
|
107
|
+
[repository](https://github.com/egordm/tenspec) holds the source.
|
|
108
|
+
|
|
109
|
+
## Attribution
|
|
110
|
+
|
|
111
|
+
The shape notation follows [jaxtyping](https://github.com/patrick-kidger/jaxtyping), which is
|
|
112
|
+
MIT licensed. jaxtyping is not a dependency of this package.
|
|
113
|
+
|
|
114
|
+
## License
|
|
115
|
+
|
|
116
|
+
MIT. See `LICENSE`.
|
tenspec-0.1.0/README.md
ADDED
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
# tenspec
|
|
2
|
+
|
|
3
|
+
Declare what an array must be, beside the value it describes, and check it at one boundary.
|
|
4
|
+
|
|
5
|
+
Tenspec moves the shape and dtype checks from the top of a function into its signature, where a
|
|
6
|
+
reader and a type checker both see them. It accepts or refuses the array the caller passed. It
|
|
7
|
+
converts nothing, moves nothing, and copies nothing, unless the declaration names a transform
|
|
8
|
+
that copies.
|
|
9
|
+
|
|
10
|
+
This is an early release, at version 0.1.0. The interface can change before version 1.0.
|
|
11
|
+
|
|
12
|
+
## Install
|
|
13
|
+
|
|
14
|
+
Install it with pip, or with uv. Either one is enough, so run the line for the tool you use:
|
|
15
|
+
|
|
16
|
+
```bash
|
|
17
|
+
pip install "tenspec[numpy]" # or: uv add "tenspec[numpy]"
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
Name the array library you need as an extra: `numpy` or `torch`. `import tenspec` loads neither
|
|
21
|
+
one. It needs Python 3.13 or newer.
|
|
22
|
+
|
|
23
|
+
## What it gives you
|
|
24
|
+
|
|
25
|
+
- **Ordinary types.** A declaration is an ordinary annotation carrying your backend's array
|
|
26
|
+
type, so a type checker reads it and untouched code still runs. The checkers, the settings
|
|
27
|
+
they need and their measured limits are on
|
|
28
|
+
[Type checkers and Ruff](https://egordm.github.io/tenspec/type-checkers.html).
|
|
29
|
+
- **Pydantic boundaries, natively.** `TensorContracts` relates the fields of your own model
|
|
30
|
+
inside one model validation.
|
|
31
|
+
- **Runtime checks are opt in.** An annotation alone performs no runtime work.
|
|
32
|
+
- **Transforms are explicit.** A value changes only when a declaration names a transform, and
|
|
33
|
+
the runtime verifies the result kept its class, shape, native dtype and device.
|
|
34
|
+
|
|
35
|
+
## The three entry points
|
|
36
|
+
|
|
37
|
+
| Entry point | Use it for |
|
|
38
|
+
|---|---|
|
|
39
|
+
| `@checked` | a function or a method, with its arguments and its return |
|
|
40
|
+
| `TensorContracts` | a Pydantic model, whose related fields share one boundary |
|
|
41
|
+
| `validate(value, annotation)` | one value, or one tuple of related values |
|
|
42
|
+
|
|
43
|
+
```python
|
|
44
|
+
from typing import Literal as Shape
|
|
45
|
+
|
|
46
|
+
import numpy as np
|
|
47
|
+
|
|
48
|
+
from tenspec import checked
|
|
49
|
+
from tenspec.numpy import Float
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
@checked
|
|
53
|
+
def weigh_columns(
|
|
54
|
+
values: Float[Shape["rows features"]], weights: Float[Shape["features"]]
|
|
55
|
+
) -> Float[Shape["rows features"]]:
|
|
56
|
+
return values * weights
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
print(weigh_columns(np.ones((3, 4)), np.array([1.0, 2.0, 3.0, 4.0])).shape)
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
`weights` must cover as many features as `values` has columns. A disagreement is a run-time
|
|
63
|
+
refusal naming the axis, both extents, and the bindings the boundary held.
|
|
64
|
+
|
|
65
|
+
No type checker compares two shapes. A shape that disagrees is a run-time refusal, never a
|
|
66
|
+
static error.
|
|
67
|
+
|
|
68
|
+
## Documentation
|
|
69
|
+
|
|
70
|
+
[The documentation site](https://egordm.github.io/tenspec/) holds the guide, the API reference
|
|
71
|
+
and two executed tutorials. The
|
|
72
|
+
[repository](https://github.com/egordm/tenspec) holds the source.
|
|
73
|
+
|
|
74
|
+
## Attribution
|
|
75
|
+
|
|
76
|
+
The shape notation follows [jaxtyping](https://github.com/patrick-kidger/jaxtyping), which is
|
|
77
|
+
MIT licensed. jaxtyping is not a dependency of this package.
|
|
78
|
+
|
|
79
|
+
## License
|
|
80
|
+
|
|
81
|
+
MIT. See `LICENSE`.
|
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
"""Forecasting: a NaN-aware scaler whose method shape the declarations must fit.
|
|
2
|
+
|
|
3
|
+
A demonstration against the accepted method shape of Cairn's ContextScale, in
|
|
4
|
+
`packages/cairn-nn/src/cairn/nn/normalization.py` at Cairn b869c260. No consumer file is
|
|
5
|
+
read at run time and no consumer file changes. The body holds only the statistics needed to
|
|
6
|
+
show that the declarations fit the real boundary.
|
|
7
|
+
|
|
8
|
+
Run it with `uv run python packages/tenspec/examples/forecast_scale.py`.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from typing import Any, Self
|
|
12
|
+
from typing import Literal as Shape
|
|
13
|
+
|
|
14
|
+
import numpy as np
|
|
15
|
+
from numpy.typing import NDArray
|
|
16
|
+
|
|
17
|
+
from pydantic import ValidationError
|
|
18
|
+
from tenspec import checked
|
|
19
|
+
from tenspec.numpy import Bool, Float, Float64
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class ContextScale:
|
|
23
|
+
"""A mean and a floored scale along the last axis, fitted once per context."""
|
|
24
|
+
|
|
25
|
+
def __init__(self, location: NDArray[np.float64], scale: NDArray[np.float64]) -> None:
|
|
26
|
+
self.location = location
|
|
27
|
+
self.scale = scale
|
|
28
|
+
|
|
29
|
+
@classmethod
|
|
30
|
+
@checked
|
|
31
|
+
def fit(
|
|
32
|
+
cls,
|
|
33
|
+
values: Float[Shape["*batch time"]],
|
|
34
|
+
mask: Bool[Shape["*batch time"]] | None = None,
|
|
35
|
+
*,
|
|
36
|
+
flatness_floor: float = 1e-3,
|
|
37
|
+
) -> Self:
|
|
38
|
+
"""Fit over the last axis. A NaN and a False mask step are both missing.
|
|
39
|
+
|
|
40
|
+
Returns:
|
|
41
|
+
A scaler of the class that received the call.
|
|
42
|
+
"""
|
|
43
|
+
observed = values if mask is None else np.where(mask, values, np.nan)
|
|
44
|
+
location = np.nanmean(observed, axis=-1)
|
|
45
|
+
spread = np.nanstd(observed, axis=-1)
|
|
46
|
+
floor = flatness_floor * (np.abs(location) + 1.0)
|
|
47
|
+
return cls(np.asarray(location, np.float64), np.asarray(np.maximum(spread, floor)))
|
|
48
|
+
|
|
49
|
+
@checked
|
|
50
|
+
def normalize(self, values: Float[Shape["*batch time"]]) -> Float64[Shape["*batch time"]]:
|
|
51
|
+
"""Centre and scale, and let a missing observation pass through.
|
|
52
|
+
|
|
53
|
+
Returns:
|
|
54
|
+
The normalized block, in float64.
|
|
55
|
+
"""
|
|
56
|
+
return (values - self.location[..., None]) / self.scale[..., None]
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class RobustScale(ContextScale):
|
|
60
|
+
"""A subclass. Its fit must return a RobustScale, never a ContextScale."""
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def main() -> None:
|
|
64
|
+
"""Fit and normalize a target history and a covariate history of different ranks."""
|
|
65
|
+
target = np.array([1.0, np.nan, 3.0, 5.0], dtype=np.float32)
|
|
66
|
+
covariates = np.array([[2.0, 4.0, 6.0, 8.0], [1.0, 1.0, np.nan, 1.0]], dtype=np.float32)
|
|
67
|
+
observed = np.array([[True, True, True, False], [True, True, True, True]])
|
|
68
|
+
|
|
69
|
+
target_scale = ContextScale.fit(target)
|
|
70
|
+
covariate_scale = ContextScale.fit(covariates, observed, flatness_floor=0.05)
|
|
71
|
+
print(f"Target location {target_scale.location.round(3).tolist()}")
|
|
72
|
+
print(f"Covariate location {covariate_scale.location.round(3).tolist()}")
|
|
73
|
+
|
|
74
|
+
normalized: Any = covariate_scale.normalize(covariates)
|
|
75
|
+
print(f"Normalized {normalized.shape}, dtype {normalized.dtype}")
|
|
76
|
+
print(f"A missing observation passes through: {bool(np.isnan(normalized[1, 2]))}")
|
|
77
|
+
|
|
78
|
+
print(f"A subclass fit returns its own class: {type(RobustScale.fit(target)).__name__}")
|
|
79
|
+
|
|
80
|
+
try:
|
|
81
|
+
ContextScale.fit(covariates, np.array([True, False, True, True]))
|
|
82
|
+
except ValidationError as refusal:
|
|
83
|
+
print(f"The mask must cover the same axes: {refusal.errors()[0]['msg']}")
|
|
84
|
+
|
|
85
|
+
try:
|
|
86
|
+
ContextScale.fit(np.array([1, 2, 3]))
|
|
87
|
+
except ValidationError as refusal:
|
|
88
|
+
print(f"An integer history is refused: {refusal.errors()[0]['msg']}")
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
if __name__ == "__main__":
|
|
92
|
+
main()
|
|
@@ -0,0 +1,161 @@
|
|
|
1
|
+
"""LOCOS: a capture boundary, and related measurement fields.
|
|
2
|
+
|
|
3
|
+
A demonstration against the real operand relationships of
|
|
4
|
+
`packages/skllm/src/skllm/attribution/locos/torch.py` and `.../locos/models.py` at Nightjar
|
|
5
|
+
476e2743. No consumer file is read at run time and no consumer file changes. The numerical
|
|
6
|
+
bodies are the smallest work that shows the declarations fit the real boundary.
|
|
7
|
+
|
|
8
|
+
Run it with `uv run python packages/tenspec/examples/locos_capture.py`.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from typing import Annotated, Any, TypeVar
|
|
12
|
+
from typing import Literal as Shape
|
|
13
|
+
|
|
14
|
+
import numpy as np
|
|
15
|
+
import torch
|
|
16
|
+
from numpy.typing import NDArray
|
|
17
|
+
from torch import Tensor
|
|
18
|
+
|
|
19
|
+
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, ValidationError
|
|
20
|
+
from tenspec import Finite, Floating, Materialized, NonEmpty, TensorContracts, checked, validate
|
|
21
|
+
from tenspec.numpy import Float64, ReadOnly
|
|
22
|
+
from tenspec.torch import Float
|
|
23
|
+
from tenspec.types.properties import Device
|
|
24
|
+
|
|
25
|
+
Scalar = TypeVar("Scalar", bound=Floating)
|
|
26
|
+
Placement = TypeVar("Placement", bound=Device)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@checked
|
|
30
|
+
def head_directions[Scalar: Floating, Placement: Device](
|
|
31
|
+
output_projection: Float[
|
|
32
|
+
Shape["heads head_dim hidden"], Scalar, Placement, NonEmpty, Materialized
|
|
33
|
+
],
|
|
34
|
+
token_direction: Float[Shape["hidden"], Scalar, Placement, NonEmpty, Materialized],
|
|
35
|
+
) -> Tensor:
|
|
36
|
+
"""One direction per head, in the operands' own format and placement.
|
|
37
|
+
|
|
38
|
+
The contraction needs a populated axis and actual values on both operands, so both
|
|
39
|
+
say so. An empty head dimension would otherwise reach the einsum and return a
|
|
40
|
+
direction of no length.
|
|
41
|
+
"""
|
|
42
|
+
return torch.einsum("hdm,m->hd", output_projection, token_direction)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def measure_step(
|
|
46
|
+
attention: Tensor, values: Tensor, directions: Tensor, source_span: tuple[int, int]
|
|
47
|
+
) -> tuple[Tensor, Tensor]:
|
|
48
|
+
"""The contribution of a source span, and the contribution outside it, per head.
|
|
49
|
+
|
|
50
|
+
The capture boundary checks the related operands together, then the numerical body runs
|
|
51
|
+
with no further check. Every operand must be populated and hold actual values, because
|
|
52
|
+
the modulo and the contraction below read both. Grouped-head divisibility and the span
|
|
53
|
+
bounds stay ordinary Python, because neither is a relation between axes.
|
|
54
|
+
|
|
55
|
+
Returns:
|
|
56
|
+
The source contribution and the background contribution, one value per head.
|
|
57
|
+
|
|
58
|
+
Raises:
|
|
59
|
+
ValidationError: when an operand is empty, or disagrees with a related operand.
|
|
60
|
+
ValueError: when the head groups do not divide, or the span leaves no background.
|
|
61
|
+
"""
|
|
62
|
+
attention, values, directions = validate(
|
|
63
|
+
(attention, values, directions),
|
|
64
|
+
tuple[
|
|
65
|
+
Float[Shape["heads keys"], Scalar, Placement, NonEmpty, Materialized],
|
|
66
|
+
Float[Shape["kv_heads keys head_dim"], Scalar, Placement, NonEmpty, Materialized],
|
|
67
|
+
Float[Shape["heads head_dim"], Scalar, Placement, NonEmpty, Materialized],
|
|
68
|
+
],
|
|
69
|
+
)
|
|
70
|
+
heads, keys = attention.shape
|
|
71
|
+
kv_heads = values.shape[0]
|
|
72
|
+
if heads % kv_heads:
|
|
73
|
+
raise ValueError(f"{heads} query heads do not divide over {kv_heads} key-value heads")
|
|
74
|
+
start, stop = source_span
|
|
75
|
+
if not 0 <= start < stop <= keys or stop - start == keys:
|
|
76
|
+
raise ValueError(f"span {source_span} leaves no background inside {keys} keys")
|
|
77
|
+
|
|
78
|
+
grouped = directions.reshape(kv_heads, heads // kv_heads, values.shape[2])
|
|
79
|
+
per_key = torch.einsum("kjd,kgd->kgj", values, grouped).reshape(heads, keys)
|
|
80
|
+
weighted = attention * per_key
|
|
81
|
+
return weighted[:, start:stop].sum(dim=1), weighted.sum(dim=1) - weighted[:, start:stop].sum(
|
|
82
|
+
dim=1
|
|
83
|
+
)
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def sealed(values: NDArray[np.float64]) -> NDArray[np.float64]:
|
|
87
|
+
"""A read-only copy, the way the consumer's own base prepares a block.
|
|
88
|
+
|
|
89
|
+
The result owns its storage. It is a copy, not a view of what the caller passed.
|
|
90
|
+
|
|
91
|
+
Returns:
|
|
92
|
+
The read-only copy.
|
|
93
|
+
"""
|
|
94
|
+
copy = np.array(values, dtype=np.float64, copy=True)
|
|
95
|
+
copy.setflags(write=False)
|
|
96
|
+
return copy
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
Block = Annotated[
|
|
100
|
+
Float64[Shape["selected_steps layers heads"], Finite, ReadOnly], BeforeValidator(sealed)
|
|
101
|
+
]
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
class FrozenArrays(BaseModel):
|
|
105
|
+
"""The consumer's own base. Tenspec adds none of its own and changes no setting."""
|
|
106
|
+
|
|
107
|
+
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
class LocosAttribution(TensorContracts, FrozenArrays):
|
|
111
|
+
"""Two representative measurement fields, on the same axes, stated once per field."""
|
|
112
|
+
|
|
113
|
+
source_attention: Block = Field(description="Source attention per step, layer, and head.")
|
|
114
|
+
background_attention: Block = Field(description="Background attention on the same axes.")
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def main() -> None:
|
|
118
|
+
"""Run the capture boundary and the measurement model."""
|
|
119
|
+
heads, kv_heads, keys, head_dim, hidden = 4, 2, 6, 3, 5
|
|
120
|
+
projection = torch.ones(heads, head_dim, hidden)
|
|
121
|
+
token = torch.ones(hidden)
|
|
122
|
+
directions = head_directions(projection, token)
|
|
123
|
+
print(f"Directions: {tuple(directions.shape)}, dtype {directions.dtype}")
|
|
124
|
+
|
|
125
|
+
attention = torch.full((heads, keys), 1.0 / keys)
|
|
126
|
+
values = torch.ones(kv_heads, keys, head_dim)
|
|
127
|
+
source, background = measure_step(attention, values, directions, (1, 3))
|
|
128
|
+
print(f"Source contribution {source.tolist()}")
|
|
129
|
+
print(f"Background contribution {background.tolist()}")
|
|
130
|
+
|
|
131
|
+
try:
|
|
132
|
+
measure_step(attention, torch.ones(3, keys, head_dim), directions, (1, 3))
|
|
133
|
+
except ValueError as refusal:
|
|
134
|
+
print(f"Grouped heads stay ordinary Python: {refusal}")
|
|
135
|
+
|
|
136
|
+
try:
|
|
137
|
+
measure_step(attention, torch.ones(kv_heads, keys + 1, head_dim), directions, (1, 3))
|
|
138
|
+
except ValidationError as refusal:
|
|
139
|
+
print(f"The shared key axis is a declaration: {refusal.errors()[0]['msg']}")
|
|
140
|
+
|
|
141
|
+
try:
|
|
142
|
+
measure_step(attention, torch.ones(0, keys, head_dim), directions, (1, 3))
|
|
143
|
+
except ValidationError as refusal:
|
|
144
|
+
print(f"An empty operand refuses before the modulo: {refusal.errors()[0]['msg']}")
|
|
145
|
+
|
|
146
|
+
measurements = LocosAttribution(
|
|
147
|
+
source_attention=np.ones((2, 3, heads)), background_attention=np.zeros((2, 3, heads))
|
|
148
|
+
)
|
|
149
|
+
stored: Any = measurements.source_attention
|
|
150
|
+
print(f"Stored block {stored.shape}, writeable {stored.flags.writeable}")
|
|
151
|
+
|
|
152
|
+
try:
|
|
153
|
+
LocosAttribution(
|
|
154
|
+
source_attention=np.ones((2, 3, heads)), background_attention=np.zeros((5, 3, heads))
|
|
155
|
+
)
|
|
156
|
+
except ValidationError as refusal:
|
|
157
|
+
print(f"Related fields share their axes: {refusal.errors()[0]['msg']}")
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
if __name__ == "__main__":
|
|
161
|
+
main()
|
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "tenspec"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Tensor specifications that a function, a model, or a direct call checks at its boundary."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.13"
|
|
7
|
+
license = "MIT"
|
|
8
|
+
license-files = ["LICENSE"]
|
|
9
|
+
authors = [{ name = "Egor Dmitriev" }]
|
|
10
|
+
classifiers = [
|
|
11
|
+
"Development Status :: 3 - Alpha",
|
|
12
|
+
"License :: OSI Approved :: MIT License",
|
|
13
|
+
"Programming Language :: Python :: 3.13",
|
|
14
|
+
"Programming Language :: Python :: 3.14",
|
|
15
|
+
"Typing :: Typed",
|
|
16
|
+
]
|
|
17
|
+
dependencies = [
|
|
18
|
+
"pydantic>=2.13.4,<3",
|
|
19
|
+
"typing-extensions>=4.13",
|
|
20
|
+
]
|
|
21
|
+
|
|
22
|
+
[project.urls]
|
|
23
|
+
Source = "https://github.com/egordm/tenspec"
|
|
24
|
+
Documentation = "https://egordm.github.io/tenspec/"
|
|
25
|
+
Issues = "https://github.com/egordm/tenspec/issues"
|
|
26
|
+
|
|
27
|
+
[project.optional-dependencies]
|
|
28
|
+
# One extra per array backend. The common tenspec import loads neither library.
|
|
29
|
+
#
|
|
30
|
+
# One floor per Python release. Each one is the oldest release tested on the interpreter it
|
|
31
|
+
# applies to: numpy 2.1 and torch 2.6 on Python 3.13, and numpy 2.3.2 and torch 2.9 on 3.14.
|
|
32
|
+
numpy = [
|
|
33
|
+
"numpy>=2.1; python_version < '3.14'",
|
|
34
|
+
"numpy>=2.3.2; python_version >= '3.14'",
|
|
35
|
+
]
|
|
36
|
+
torch = [
|
|
37
|
+
"torch>=2.6; python_version < '3.14'",
|
|
38
|
+
"torch>=2.9; python_version >= '3.14'",
|
|
39
|
+
]
|
|
40
|
+
dev = [
|
|
41
|
+
"tenspec[numpy,torch]",
|
|
42
|
+
"pytest>=8.0",
|
|
43
|
+
"ruff>=0.8",
|
|
44
|
+
"ty>=0.0.78",
|
|
45
|
+
"poethepoet>=0.30",
|
|
46
|
+
]
|
|
47
|
+
|
|
48
|
+
[build-system]
|
|
49
|
+
requires = ["hatchling"]
|
|
50
|
+
build-backend = "hatchling.build"
|
|
51
|
+
|
|
52
|
+
[tool.hatch.build.targets.wheel]
|
|
53
|
+
packages = ["src/tenspec"]
|
|
54
|
+
|
|
55
|
+
[tool.hatch.build.targets.sdist]
|
|
56
|
+
# README.md joins the list because `readme` names it: a source archive must be able to rebuild
|
|
57
|
+
# its own metadata.
|
|
58
|
+
include = ["src", "tests", "examples", "pyproject.toml", "README.md"]
|
|
59
|
+
|
|
60
|
+
# ruff + ty + pytest + poe are configured ONCE at the workspace root (../../pyproject.toml,
|
|
61
|
+
# ruff.toml, poe_tasks.toml): this member carries none of its own.
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""Tenspec: declare what an array must be, and check it at one boundary.
|
|
2
|
+
|
|
3
|
+
Write the declaration once, beside the value it describes. A checked function, a model
|
|
4
|
+
mixin, or a standalone call then holds it. Import a backend module for the array library
|
|
5
|
+
you use: `tenspec.numpy` or `tenspec.torch`. This module loads neither one.
|
|
6
|
+
|
|
7
|
+
This is the common surface. Everything below it keeps its defining module, including
|
|
8
|
+
`tenspec.errors.TensorMismatch`, `tenspec.errors.BindingError`,
|
|
9
|
+
`tenspec.runtime.validation.validate_array`, `tenspec.runtime.bindings.Bindings`,
|
|
10
|
+
`tenspec.runtime.arrays.ArrayBackend`, and `tenspec.pydantic.annotations.prepare_annotation`.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from typing import Literal as Shape
|
|
14
|
+
|
|
15
|
+
from tenspec.errors import AnnotationError, ConstraintEvaluationError
|
|
16
|
+
from tenspec.pydantic.calls import checked
|
|
17
|
+
from tenspec.pydantic.models import TensorContracts
|
|
18
|
+
from tenspec.pydantic.standalone import validate
|
|
19
|
+
from tenspec.runtime.bindings import axis_size
|
|
20
|
+
from tenspec.types.dtypes import DType, Floating
|
|
21
|
+
from tenspec.types.properties import Contiguous, Finite, Materialized, NonEmpty
|
|
22
|
+
|
|
23
|
+
__version__ = "0.1.0"
|
|
24
|
+
|
|
25
|
+
__all__ = [
|
|
26
|
+
"AnnotationError",
|
|
27
|
+
"ConstraintEvaluationError",
|
|
28
|
+
"Contiguous",
|
|
29
|
+
"DType",
|
|
30
|
+
"Finite",
|
|
31
|
+
"Floating",
|
|
32
|
+
"Materialized",
|
|
33
|
+
"NonEmpty",
|
|
34
|
+
"Shape",
|
|
35
|
+
"TensorContracts",
|
|
36
|
+
"axis_size",
|
|
37
|
+
"checked",
|
|
38
|
+
"validate",
|
|
39
|
+
]
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
"""The errors Tenspec raises, and what each one reports."""
|
|
2
|
+
|
|
3
|
+
from collections.abc import Mapping
|
|
4
|
+
from enum import StrEnum
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class AnnotationError(TypeError):
|
|
8
|
+
"""A declaration is invalid, unsupported, or never reached Tenspec's preparation."""
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class ConstraintEvaluationError(RuntimeError):
|
|
12
|
+
"""A backend cannot inspect a declared property of this value.
|
|
13
|
+
|
|
14
|
+
This never reports a measured mismatch. It reports that the check could not run.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class BindingReason(StrEnum):
|
|
19
|
+
"""Why a dimension lookup found no extent."""
|
|
20
|
+
|
|
21
|
+
NO_SCOPE = "no_scope"
|
|
22
|
+
UNBOUND = "unbound"
|
|
23
|
+
NOT_SCALAR = "not_scalar"
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class BindingError(ValueError):
|
|
27
|
+
"""A required dimension was unavailable, rather than different.
|
|
28
|
+
|
|
29
|
+
The reason says which lookup failed. The name is the dimension the caller asked for.
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
def __init__(self, reason: BindingReason, name: str | None = None) -> None:
|
|
33
|
+
self.reason = reason
|
|
34
|
+
self.name = name
|
|
35
|
+
messages = {
|
|
36
|
+
BindingReason.NO_SCOPE: "no Tenspec validation scope is open",
|
|
37
|
+
BindingReason.UNBOUND: f"dimension {name!r} has no binding in this scope",
|
|
38
|
+
BindingReason.NOT_SCALAR: f"dimension {name!r} holds a group of axes, not one extent",
|
|
39
|
+
}
|
|
40
|
+
super().__init__(messages[reason])
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class TensorMismatch(ValueError): # noqa: N818 the accepted contract fixes this name.
|
|
44
|
+
"""An array disagrees with its declared requirement.
|
|
45
|
+
|
|
46
|
+
The requirement, the observed fact, and the relevant bindings stay readable on the
|
|
47
|
+
error. The error holds no array.
|
|
48
|
+
"""
|
|
49
|
+
|
|
50
|
+
def __init__(
|
|
51
|
+
self,
|
|
52
|
+
*,
|
|
53
|
+
requirement: str,
|
|
54
|
+
observed: str,
|
|
55
|
+
bindings: Mapping[str, int | tuple[int, ...]] | None = None,
|
|
56
|
+
) -> None:
|
|
57
|
+
self.requirement = requirement
|
|
58
|
+
self.observed = observed
|
|
59
|
+
self.bindings: Mapping[str, int | tuple[int, ...]] = dict(bindings or {})
|
|
60
|
+
detail = f", with {self.bindings}" if self.bindings else ""
|
|
61
|
+
super().__init__(f"expected {requirement}, but the array has {observed}{detail}")
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def describe_empty(shape: tuple[int, ...]) -> str:
|
|
65
|
+
"""What a refusal reports about an array that holds no elements.
|
|
66
|
+
|
|
67
|
+
It reads the shape alone. Axis indices start at 0.
|
|
68
|
+
|
|
69
|
+
describe_empty((2, 0, 0)) returns "shape (2, 0, 0), with no elements at axis 1, axis 2".
|
|
70
|
+
|
|
71
|
+
Returns:
|
|
72
|
+
The observed text for a TensorMismatch.
|
|
73
|
+
"""
|
|
74
|
+
empty = ", ".join(f"axis {index}" for index, extent in enumerate(shape) if extent == 0)
|
|
75
|
+
return f"shape {shape}, with no elements at {empty}"
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
class TransformContractError(RuntimeError):
|
|
79
|
+
"""A transform changed a structural fact that it had to preserve.
|
|
80
|
+
|
|
81
|
+
This reports a defect in the transform, not invalid caller input, so it is a
|
|
82
|
+
RuntimeError and never a ValueError. It names the transform and the one fact that
|
|
83
|
+
changed. Every field is text, so the error holds no array.
|
|
84
|
+
"""
|
|
85
|
+
|
|
86
|
+
def __init__(self, *, transform: str, fact: str, before: str, after: str) -> None:
|
|
87
|
+
self.transform = transform
|
|
88
|
+
self.fact = fact
|
|
89
|
+
super().__init__(f"{transform} changed the {fact} of its input, from {before} to {after}")
|