stackformers 3.8.0__tar.gz

stackformers-3.8.0/.claude/settings.json

@@ -0,0 +1,7 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(python -c \"import torch; print\\(torch.__version__\\)\")"
+    ]
+  }
+}

stackformers-3.8.0/.claudeignore

@@ -0,0 +1,17 @@
+# Lock files - never useful to read
+uv.lock
+poetry.lock
+package-lock.json
+yarn.lock
+Cargo.lock
+*.lock
+
+# Other noise
+__pycache__/
+*.pyc
+.venv/
+dist/
+build/
+*.egg-info/
+.mypy_cache/
+.ruff_cache/

stackformers-3.8.0/.gitignore

@@ -0,0 +1,219 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+# Temporary file for partial code execution
+tempCodeRunnerFile.py
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+x-transformers

stackformers-3.8.0/.python-version

@@ -0,0 +1 @@
+3.11

stackformers-3.8.0/.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python"
+}

stackformers-3.8.0/CLAUDE.md

@@ -0,0 +1,89 @@
+# CLAUDE.md — stackformers
+
+## What this project is
+Typed, composable, SOLID transformer library for PyTorch.
+Behavior comes from injected dependencies, not constructor flags.
+
+## Key rules — never violate these
+
+- No optional fields on dataclasses to encode different modes — use sealed unions instead
+- No `if self.x is not None` inside `forward()` — use null objects (e.g. `NoPosEncoding`)
+- No flags in `__init__` that change `forward()` behavior — new behavior = new class
+- `forward()` must be traceable by torch.compile and torch.export — no Python control flow on tensors
+- One concept per file — if you can't describe it in one sentence, split it
+- All tensor args annotated with jaxtyping shape strings using project dim vocabulary
+- Every collaborator accepted by `__init__` must be typed as a protocol, never as a concrete class
+
+## Dim naming convention (use everywhere)
+| Symbol | Meaning |
+|--------|---------|
+| b      | batch size |
+| n      | query sequence length |
+| s      | key/value sequence length (= n for self-attn) |
+| d      | model dim |
+| h      | number of heads |
+| dh     | dim per head |
+| w      | window size |
+| nt     | total tokens in packed sequence |
+
+## File organisation
+
+- All code lives directly under `stackformers/`. Breaking changes go in a new top-level package.
+- Configs live next to the class they configure: `attention/config.py`, `feedforward/config.py`.
+- Cross-cutting configs (`LayerConfig`, `EncoderConfig`, `DecoderConfig`) live at `config.py`.
+- Each module has a `factory.py` with a `build_*` function that dispatches on config type via `match`.
+- Opinionated presets live in `presets/` — one file per preset class.
+- Protocols are defined in the module they belong to (`attention/protocols.py`, etc.), not a central file.
+- Implementations satisfy protocols structurally — they never import the protocol they implement.
+
+## Config and factory conventions
+
+Every union config type uses a `kind: Literal[...]` discriminator field so configs round-trip through JSON unambiguously. The union itself is an `Annotated[..., Field(discriminator="kind")]` alias defined in the same `config.py`.
+
+Each module owns its builder in a co-located `factory.py`. The function signature is `build_*(config: *Config, ...) -> Protocol`. Presets call these factories — they never instantiate concrete classes directly.
+
+## Module-level READMEs
+
+Each subdirectory has a `README.md`. When you change a module's design — add a variant, remove a concept, change a dispatch strategy — update its `README.md` in the same commit. The README covers: purpose in one or two sentences, key design decisions, and how to extend with a new variant. It must not restate what reading the source already gives you (no file tables, no copied signatures).
+
+## Testing rules
+- One test file per source file, mirroring directory structure
+- Use pytest fixtures for module construction, inputs, and sequence objects — no plain helper functions
+- Parametrise over `device_dtype` (all device×dtype combos) for compute tests; use `device` alone (float32) for gradient tests
+- Test shape contracts, not numeric values
+- Numerical assertions use `atol(dtype)` from `tests/conftest.py`
+
+## Dependencies
+- PyTorch native first: `F.scaled_dot_product_attention`, `torch.nn.attention.varlen.varlen_attn`
+- No optional third-party kernel dependencies — all kernels are pure PyTorch
+- Pydantic for configs; use `Field(gt=0)` / `Field(ge=0)` constraints, not `@field_validator`, for simple bounds
+- Prefer `nn.Module` activations (`nn.SiLU()`, `nn.GELU()`) over `torch.nn.functional` calls
+
+## x-transformers reference
+Cloned at `./x-transformers/` — read for math and implementation details only.
+Do not copy its architecture.
+
+## Versioning
+
+Version is maintained in `pyproject.toml` under `[project].version` and exposed at runtime via `stackformers.__version__`.
+
+Use [Semantic Versioning](https://semver.org/):
+- **MAJOR** — breaking public API change
+- **MINOR** — new backwards-compatible feature
+- **PATCH** — bug fix or internal change with no API impact
+
+Bump the version in `pyproject.toml` on every commit that changes behaviour. Commits that only touch docs, tests, or CI do not require a bump.
+
+## Commands
+```bash
+just fmt          # ruff format
+just lint         # ruff check
+just types        # pyrefly
+just test         # pytest
+just check        # full CI gate: fmt-check + lint + types + test
+```
+
+## PyTorch
+- Before writing any torch code, run: `python -c "import torch; print(torch.__version__)"`
+- Always prefer built-in torch implementations over custom ones
+- Check `.venv/lib/python3.*/site-packages/torch/nn/modules/` for what's available

stackformers-3.8.0/Justfile

@@ -0,0 +1,68 @@
+set dotenv-load := false
+
+# List available recipes
+default:
+    @just --list
+
+# ── Dependencies ────────────────────────────────────────────────────────────
+
+# Install all dependencies (including dev)
+sync:
+    uv sync --group dev
+
+# ── Formatting ──────────────────────────────────────────────────────────────
+
+# Format source and tests
+fmt:
+    uv run ruff format stackformers/ tests/
+
+# Check formatting without modifying files
+fmt-check:
+    uv run ruff format --check stackformers/ tests/
+
+# ── Linting ─────────────────────────────────────────────────────────────────
+
+# Lint source and tests
+lint:
+    uv run ruff check stackformers/ tests/
+
+# Lint and auto-fix what's safe
+lint-fix:
+    uv run ruff check --fix stackformers/ tests/
+
+# ── Type checking ────────────────────────────────────────────────────────────
+
+# Type-check with pyrefly
+types:
+    uv run pyrefly check stackformers/ tests/
+
+# ── Testing ─────────────────────────────────────────────────────────────────
+
+# Run tests
+test:
+    uv run pytest tests/ -q
+
+# Run tests with coverage report
+test-cov:
+    uv run pytest tests/ --cov=stackformers --cov-report=term-missing -q
+
+# Run a single test file or pattern, e.g.: just test-only tests/v1/test_encoder.py
+test-only target:
+    uv run pytest {{ target }} -v
+
+# ── Combined checks ──────────────────────────────────────────────────────────
+
+# Run fmt-check + lint + types + test (CI gate)
+check: fmt-check lint types test
+
+# Format, lint-fix, then run all checks
+fix: fmt lint-fix
+    just check
+
+# ── Housekeeping ─────────────────────────────────────────────────────────────
+
+# Remove build artifacts and caches
+clean:
+    rm -rf dist/ .venv/ .ruff_cache/ .pytest_cache/ __pycache__
+    find . -type d -name "__pycache__" -exec rm -rf {} +
+    find . -type d -name "*.egg-info" -exec rm -rf {} +

stackformers-3.8.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Vadym Stupakov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

stackformers-3.8.0/PKG-INFO

@@ -0,0 +1,204 @@
+Metadata-Version: 2.4
+Name: stackformers
+Version: 3.8.0
+Summary: Typed, composable, SOLID transformer library for PyTorch
+Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
+Project-URL: Repository, https://github.com/Red-Eyed/stackformers
+Project-URL: Bug Tracker, https://github.com/Red-Eyed/stackformers/issues
+Author-email: Vadym Stupakov <vadim.stupakov@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2026 Vadym Stupakov
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: attention,deep-learning,pytorch,transformer
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: beartype>=0.18
+Requires-Dist: einops>=0.7
+Requires-Dist: jaxtyping>=0.2
+Requires-Dist: numpy>=2.4.4
+Requires-Dist: pydantic>=2.0
+Requires-Dist: torch>=2.2
+Description-Content-Type: text/markdown
+
+# stackformers
+
+Typed, composable transformer library for PyTorch.
+Every architectural choice — positional encoding, normalization, feedforward variant — is an injected dependency, not a constructor flag.
+
+```bash
+uv add stackformers
+```
+
+---
+
+## Why
+
+Most transformer libraries grow into a tangle of `if self.use_rope`, `if self.window_size is not None`, and god-config objects with thirty nullable fields. Adding a new variant means touching existing code.
+
+stackformers takes a different approach:
+
+- **Swap any component without touching anything else** — `SelfAttention(config, pos_encoding=RoPE)` vs `SelfAttention(config, pos_encoding=ALiBi)` — same call site, different object
+- **No `None` checks in `forward()`** — `NoPosEncoding` is a real object that passes q/k unchanged; the branch never exists
+- **Sealed sequence unions** — `PaddedInput | PackedInput` instead of optional `cu_seqlens` and `mask` arguments that conflict with each other
+- **`torch.compile` / `torch.export` safe** — no Python control flow on tensors inside any `forward()`
+- **Structural protocols** — bring your own implementation; no ABC inheritance required
+
+---
+
+## Quick start
+
+### Zero boilerplate
+
+```python
+import torch
+from stackformers import TransformerEncoder, plain_encoder_config, make_padded_input
+
+model = TransformerEncoder(plain_encoder_config(dim=512, heads=8, num_layers=6))
+
+x    = torch.randn(2, 128, 512)
+mask = torch.ones(2, 128, dtype=torch.bool)
+out  = model(make_padded_input(x, mask))   # (2, 128, 512)
+```
+
+Switch to packed (variable-length, no padding waste) — same weights:
+
+```python
+from stackformers import make_packed_input
+
+cu  = torch.tensor([0, 64, 128], dtype=torch.int32)
+out = model(make_packed_input(x_flat, cu, max_seqlen=64))  # (128, 512)
+```
+
+Causal LM backbone:
+
+```python
+plain_encoder_config(dim=768, heads=12, num_layers=12, causal=True)
+```
+
+Sliding-window local attention (O(n · w)):
+
+```python
+from stackformers import windowed_encoder_config
+windowed_encoder_config(dim=512, heads=8, num_layers=6, window_size=128)
+```
+
+Encoder–decoder:
+
+```python
+from stackformers import TransformerDecoder, plain_decoder_config
+
+model = TransformerDecoder(plain_decoder_config(dim=512, heads=8, num_layers=6))
+out   = model(make_padded_input(x, mask), make_padded_input(context, ctx_mask))
+```
+
+### Explicit config
+
+Full control with JSON round-trip via `kind` discriminators:
+
+```python
+from stackformers import (
+    TransformerEncoderConfig, TransformerEncoder,
+    SelfAttentionConfig, SwiGLUConfig, RMSNormConfig, RoPE1DConfig,
+    make_padded_input,
+)
+
+cfg = TransformerEncoderConfig(
+    attn=SelfAttentionConfig(dim=512, heads=8, dim_head=64, causal=False),
+    ff=SwiGLUConfig(dim=512, mult=4.0),
+    norm=RMSNormConfig(dim=512),
+    pos_encoding=RoPE1DConfig(dim_head=64),
+    num_layers=6,
+)
+model = TransformerEncoder(cfg)
+
+# Serialise / restore
+cfg2 = TransformerEncoderConfig.model_validate(cfg.model_dump())
+```
+
+### Custom wiring
+
+Wire layers yourself when presets aren't enough:
+
+```python
+from stackformers import (
+    SelfAttention, SwiGLU, TransformerLayer, Encoder, RMSNorm,
+    RotaryEmbedding1D,
+    SelfAttentionConfig, SwiGLUConfig, RMSNormConfig, RoPE1DConfig,
+)
+
+pos  = RotaryEmbedding1D(RoPE1DConfig(dim_head=64))
+attn = SelfAttention(SelfAttentionConfig(dim=512, heads=8, dim_head=64), pos_encoding=pos)
+
+layers = [
+    TransformerLayer(
+        self_attn=attn,
+        ff=SwiGLU(SwiGLUConfig(dim=512)),
+        norm_attn=RMSNorm(RMSNormConfig(dim=512)),
+        norm_ff=RMSNorm(RMSNormConfig(dim=512)),
+    )
+    for _ in range(6)
+]
+encoder = Encoder(layers=layers, final_norm=RMSNorm(RMSNormConfig(dim=512)))
+```
+
+---
+
+## What's included
+
+| Area | Variants |
+|------|----------|
+| Self-attention | Global, sliding-window (local); padded and packed backends; GQA / MQA |
+| Cross-attention | Global; padded and packed backends |
+| Positional encoding | RoPE-1D, RoPE-2D, none (null object) |
+| Feedforward | SwiGLU, GEGLU |
+| Normalization | RMSNorm, LayerNorm |
+| Presets | Encoder, Decoder, CrossAttender |
+
+On CUDA with fp16/bf16 the packed path uses `torch.nn.attention.varlen.varlen_attn`. CPU and fp32 fall back to a scatter-to-padded SDPA — correct everywhere, fast where it matters.
+
+---
+
+## Development
+
+```bash
+git clone <repo> && cd stackformers
+uv sync --group dev
+
+just fmt      # format
+just lint     # lint
+just types    # type-check
+just test     # test
+just check    # full CI gate
+```
+
+---
+
+## License
+
+See [LICENSE](LICENSE).