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.
Files changed (48) hide show
  1. tenspec-0.1.0/.gitignore +22 -0
  2. tenspec-0.1.0/LICENSE +21 -0
  3. tenspec-0.1.0/PKG-INFO +116 -0
  4. tenspec-0.1.0/README.md +81 -0
  5. tenspec-0.1.0/examples/forecast_scale.py +92 -0
  6. tenspec-0.1.0/examples/locos_capture.py +161 -0
  7. tenspec-0.1.0/pyproject.toml +61 -0
  8. tenspec-0.1.0/src/tenspec/__init__.py +39 -0
  9. tenspec-0.1.0/src/tenspec/errors.py +89 -0
  10. tenspec-0.1.0/src/tenspec/numpy.py +343 -0
  11. tenspec-0.1.0/src/tenspec/py.typed +0 -0
  12. tenspec-0.1.0/src/tenspec/pydantic/__init__.py +4 -0
  13. tenspec-0.1.0/src/tenspec/pydantic/annotations.py +407 -0
  14. tenspec-0.1.0/src/tenspec/pydantic/calls.py +136 -0
  15. tenspec-0.1.0/src/tenspec/pydantic/models.py +82 -0
  16. tenspec-0.1.0/src/tenspec/pydantic/standalone.py +99 -0
  17. tenspec-0.1.0/src/tenspec/runtime/__init__.py +1 -0
  18. tenspec-0.1.0/src/tenspec/runtime/arrays.py +313 -0
  19. tenspec-0.1.0/src/tenspec/runtime/bindings.py +168 -0
  20. tenspec-0.1.0/src/tenspec/runtime/validation.py +368 -0
  21. tenspec-0.1.0/src/tenspec/torch.py +276 -0
  22. tenspec-0.1.0/src/tenspec/types/__init__.py +1 -0
  23. tenspec-0.1.0/src/tenspec/types/base.py +13 -0
  24. tenspec-0.1.0/src/tenspec/types/dtypes.py +140 -0
  25. tenspec-0.1.0/src/tenspec/types/operations.py +114 -0
  26. tenspec-0.1.0/src/tenspec/types/properties.py +74 -0
  27. tenspec-0.1.0/src/tenspec/types/shapes.py +289 -0
  28. tenspec-0.1.0/src/tenspec/types/tensor.py +28 -0
  29. tenspec-0.1.0/tests/integration/checked_caller.py +351 -0
  30. tenspec-0.1.0/tests/integration/pyrightconfig.json +6 -0
  31. tenspec-0.1.0/tests/integration/test_checkers.py +72 -0
  32. tenspec-0.1.0/tests/pydantic/test_annotations.py +294 -0
  33. tenspec-0.1.0/tests/pydantic/test_calls.py +229 -0
  34. tenspec-0.1.0/tests/pydantic/test_future_annotations.py +118 -0
  35. tenspec-0.1.0/tests/pydantic/test_models.py +211 -0
  36. tenspec-0.1.0/tests/pydantic/test_named_aliases.py +188 -0
  37. tenspec-0.1.0/tests/pydantic/test_standalone.py +183 -0
  38. tenspec-0.1.0/tests/runtime/test_arrays.py +143 -0
  39. tenspec-0.1.0/tests/runtime/test_bindings.py +117 -0
  40. tenspec-0.1.0/tests/runtime/test_validation.py +567 -0
  41. tenspec-0.1.0/tests/test_caller_example.py +366 -0
  42. tenspec-0.1.0/tests/test_imports.py +53 -0
  43. tenspec-0.1.0/tests/test_numpy.py +164 -0
  44. tenspec-0.1.0/tests/test_torch.py +170 -0
  45. tenspec-0.1.0/tests/types/test_dtypes.py +36 -0
  46. tenspec-0.1.0/tests/types/test_shapes.py +89 -0
  47. tenspec-0.1.0/tests/types/test_shared_grammar.py +72 -0
  48. tenspec-0.1.0/tests/types/test_tensor.py +39 -0
@@ -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`.
@@ -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}")