stackformers 4.3.0__tar.gz → 5.0.0rc3__tar.gz

{stackformers-4.3.0 → stackformers-5.0.0rc3}/.gitignore

@@ -37,6 +37,7 @@ pip-log.txt
 pip-delete-this-directory.txt
 
 # Unit test / coverage reports
+/work_dir/
 htmlcov/
 .tox/
 .nox/
@@ -216,4 +217,4 @@ __marimo__/
 
 # Streamlit
 .streamlit/secrets.toml
-x-transformers
+x-transformers

stackformers-5.0.0rc3/AGENTS.md

@@ -0,0 +1,19 @@
+# Repository guidance
+
+## Repository commands
+
+Read `Justfile` before running setup, build, formatting, linting, type-checking, testing, or
+cleanup commands. Prefer its existing recipes so local and automated workflows use the same
+tool arguments and defaults.
+
+## Python structured records
+
+Never use `dataclasses` in this repository. PyTorch recognizes `NamedTuple` and `TypedDict`
+reliably across pytree, export, and compilation boundaries, while dataclasses require additional
+registration and handling.
+
+- Use `NamedTuple` for fixed-shape records with named fields, especially tensor inputs and outputs.
+- Use `TypedDict` when the value must retain mapping semantics.
+- Do not introduce `dataclass`, `@dataclass`, or imports from `dataclasses`.
+
+Ruff enforces the import restriction through `TID251`.

{stackformers-4.3.0 → stackformers-5.0.0rc3}/CHANGELOG.md

@@ -6,6 +6,252 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and
 adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
 MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
 
+## [Unreleased]
+
+## [5.0.0rc3] — 2026-09-26
+
+### New Features
+
+- Accept shared or per-layer `dim_heads`, `heads`, `causal`, `ff_mult`, `dropout`, and
+  `norm_placement` in the variable-width encoder factory. Explicit `heads` allows independent
+  model and attention widths; omitting it retains exact `d_model // dim_head` inference.
+
+### Bug Fixes
+
+- Allow model widths that are not divisible by the query-head count in self- and
+  cross-attention configs. Internal query width remains `heads * dim_head`; grouped
+  attention still requires `heads` divisible by `kv_heads`.
+
+## [5.0.0rc1] — 2026-09-26
+
+### Backwards Incompatible Changes
+
+#### Immutable attention configuration
+
+`SelfAttentionConfig` and `CrossAttentionConfig` previously allowed field assignment after
+construction. They are now immutable, preventing a shared config from reinterpreting an
+existing model's projections. Construct a new config and model for different architecture
+settings instead of assigning config fields.
+
+#### Matching cross-attention layouts
+
+Cross-attention previously accepted mixed padded/packed layouts in its annotations and failed
+during tensor reshaping. Its type contract now requires matching layouts, and mixed layouts
+raise `ValueError` at runtime. Narrow both inputs together when calling through the `CrossAttn`
+protocol, or convert them to the same layout before calling.
+
+#### Empty reconstruction selections
+
+Custom reconstruction heads must now return a differentiable zero loss for zero selected
+tokens; earlier head contracts left empty-input behavior unspecified. Add an empty-input guard
+that returns `prediction.sum()` to custom heads. Independent sampling is unchanged, so callers
+must continue to support batches in which no tokens are selected.
+
+### Bug Fixes
+
+- Return finite zero losses from both built-in MLM heads when no tokens are selected.
+- Exclude masked keys completely and return zero attention contributions and gradients when
+  padding, windows, or additive bias exclude every key, including exported ONNX execution.
+- Fix ONNX Runtime attention-mask shape errors for opsets 23–25 by explicitly expanding
+  the query dimension, including dynamic batch and sequence lengths.
+
+### Performance
+
+- Avoid repadding input features on the packed self-attention path when using `NoAttnBias`.
+
+### Developers
+
+- Include the `py.typed` marker so type checkers recognize annotations in installed wheels.
+- Use shared `returns` Result values for core layout, geometry, cache, and backend checks.
+  Preserve public tensor/configuration returns, validation exceptions, and fallback warnings;
+  retain original backend failure detail for callers that handle outcomes directly.
+- Apply rustic-python's strict Pyrefly and Ruff profile to the library, tests, and examples;
+  require Pyrefly 1.3.1 and Ruff 0.16.9 or newer and check Python 3.11 compatibility.
+- Expose typed module-call signatures without replacing PyTorch's runtime hook dispatch,
+  type the experimental attention-kernel boundary, and replace stale type suppressions.
+- Require ONNX Runtime parity at opsets 23–25 instead of accepting opset-23 failures or
+  excluding newer dynamo exports using PyTorch's legacy exporter limit.
+
+## [4.7.0b6] — 2026-08-31
+
+### Documentation
+
+- Add runnable, CPU-friendly examples for equivalent padded and packed batches, causal language
+  modeling, image patches with RoPE-2D, continuous point coordinates with RoPE-ND, and both
+  focused and fully user-owned encoder subclassing.
+
+## [4.7.0b5] — 2026-08-28
+
+### Highlights
+
+- Add an export-oriented decoder path that reuses cross-attention K/V projections and grows a
+  self-attention K/V cache one token at a time, while leaving existing encoder, decoder, and
+  attention APIs unchanged.
+
+### New Features
+
+- Add composition-only `CachedSelfAttentionWrapper`, `CachedCrossAttentionWrapper`,
+  `DecoderCrossAttentionCacheBuilder`, and `CachedDecoderWrapper` APIs. The wrappers share the
+  existing model's parameters and support every built-in positional encoding and decoder
+  normalization topology.
+- Define a C++/ONNX Runtime-friendly tensor boundary with required, non-optional cross and self
+  caches, a required `step_i` tensor, dynamically shaped batch/source/history axes, and a returned
+  self cache for the next decoder session run. Cached self-attention currently requires global
+  causal attention with `NoAttnBias`.
+
+### Documentation
+
+- Add an encoder-decoder example that exports separate encoder, cross-cache builder, and cached
+  one-token decoder ONNX graphs, then verifies their autoregressive outputs with ONNX Runtime.
+- Document cache layouts, tensor names, ownership, lifecycle, dynamic axes, and the initial
+  zero-length self-cache contract for C++ integration.
+
+### Developers
+
+- Add eager parity, gradient, export, ONNX Runtime, positional-encoding, and stable `4.6.0` public
+  API regression coverage for the cached path.
+- Standardize export-facing structured records on `NamedTuple` or `TypedDict` and enforce the
+  repository's no-dataclass policy with Ruff.
+
+## [4.7.0b4] — 2026-08-19
+
+### Backwards Incompatible Changes
+
+#### Variable-width encoder configuration
+
+`VariableWidthTransformerEncoderConfig` again stores `layers: list[VariableWidthEncoderLayerConfig]`
+instead of the `d_models`/`dim_heads` representation introduced in `4.7.0b3`. Each layer record
+keeps its attention, feed-forward, norm, positional-encoding, and attention-bias choices explicit
+and validates their dimensions together. Migrate direct `4.7.0b3` constructors to complete layer
+records. The convenience config factory now requires explicit keyword arguments; replace
+`variable_width_encoder_config(d_models, dim_heads)` with
+`variable_width_encoder_config(d_models=d_models, dim_heads=dim_heads)`.
+
+### New Features
+
+- Export all concrete feed-forward configs and implementations, learned/ND positional encodings,
+  attention-bias configs, `VariableWidthEncoderLayerConfig`, and `node_encoder_config` from the
+  package root. Variable-width layers can select any public component config without subclassing.
+
+### Developers
+
+- Add static and dynamic PyTorch/ONNX export coverage with ONNX Runtime parity for every public
+  feed-forward variant and the existing Transformer configuration paths.
+
+## [4.7.0b3] — 2026-08-19
+
+### Backwards Incompatible Changes
+
+#### Variable-width encoder configuration
+
+`VariableWidthTransformerEncoderConfig` now represents variable block geometry as parallel
+`d_models` and `dim_heads` arrays, with `causal`, `ff_mult`, `dropout`, and `norm_placement` stored
+once as shared settings. Earlier beta configs stored an array of
+`VariableWidthEncoderLayerConfig` objects under `layers`; that class and its public export are
+removed. Explicit constructors and serialized beta configs must migrate from values such as
+`config.layers[index].attn.dim` to `config.d_models[index]`. The preset continues to construct
+global attention with RoPE, RMSNorm, SwiGLU, and no attention bias. Callers of
+`variable_width_encoder_config` are unchanged.
+
+## [4.7.0b2] — 2026-08-19
+
+### Added
+
+- **`HardSwishGLU`** — a bias-free, parameter-matched gated feed-forward variant using
+  `nn.Hardswish` instead of SwiGLU's `nn.SiLU`. Select it with
+  `HardSwishGLUConfig(kind="hardswish_glu")`; serialized configurations preserve the distinction
+  between the exact SwiGLU operation and its piecewise-linear HardSwish alternative.
+
+### Compatibility
+
+- Existing feed-forward classes, config discriminators, presets, and defaults are unchanged.
+  `HardSwishGLU` is opt-in and uses the same input/output shape and hidden-width rule as SwiGLU.
+
+## [4.7.0b1] — 2026-08-18
+
+### Added
+
+- **Variable-width Transformer encoder preset** through `VariableWidthTransformerEncoder` and
+  `variable_width_encoder_config`. Callers provide one model width and head dimension per block;
+  the preset derives head counts and inserts a bias-free learned projection only where adjacent
+  model widths differ. Expanded layer configs retain explicit attention, feed-forward, norm,
+  positional-encoding, and attention-bias collaborators and support JSON round trips.
+
+### Compatibility
+
+- Existing uniform-width encoder classes, configs, serialized payloads, checkpoints, and defaults
+  are unchanged. Variable-width models use their own preset and state-dict structure.
+
+## [4.6.0] — 2026-08-14
+
+### Added
+
+- **Configurable normalization placement** for encoder, decoder, and cross-attender presets through
+  `NormPlacement = Literal["pre", "post", "sandwich", "reordered"]`. Each topology has a focused
+  layer class while the existing `TransformerLayer`, `DecoderLayer`, and `CrossAttenderLayer`
+  remain the pre-norm implementations. `"reordered"` implements the OLMo 2 residual layout
+  (`x + norm(branch(x))`), while sandwich uses distinct norms before and after every residual
+  branch—four per encoder/cross-attender layer and six per decoder layer.
+
+### Compatibility
+
+- **The default remains `"pre"`.** Omitting `norm_placement` preserves the previous operation
+  order, exact outputs and gradients, positional `TransformerLayer` constructor, parameter names,
+  and state-dict shapes across encoder, decoder, and cross-attender layers. Config payloads without
+  the new field deserialize as pre-norm. The three legacy layer classes retain their original
+  constructors and module attributes.
+- Serialized preset-config `model_dump()` output now includes
+  `"norm_placement": "pre"` unless callers request default-value exclusion. Sandwich is the only
+  mode that adds parameters because its post-branch norms must be independent.
+
+## [4.5.0] — 2026-08-14
+
+### Added
+
+- **`GELUFFN`** — a standard bias-free, two-projection Transformer feed-forward network using
+  `nn.GELU(approximate="tanh")`. Select it with `GELUConfig(kind="gelu")`; its hidden width is
+  `int(dim * mult)`, matching other non-gated FFNs.
+
+## [4.4.0] — 2026-07-17
+
+### Added
+
+- **`CosineHead`** (`stackformers/mlm/head_cosine.py`) — a `ReconstructionHead` that scores
+  a linear projection of the encoder's masked-position output against the target via
+  `1 - cosine_similarity`, instead of `RegressionHead`'s MSE.
+
+### Changed
+
+- **`MLMWrapper`'s default `head` is now `CosineHead`, not `RegressionHead`.**
+  `mlm_loss`'s target is `input.x`, and nothing constrains its scale: once the tokenizer
+  trains purely under the main task (`mlm_loss` has zero gradient into it, per 4.3.1),
+  its output distribution is free to drift arbitrarily over training. MSE is sensitive
+  to that drift — matching magnitude is part of what it scores — while cosine
+  similarity scores direction only, so it stays meaningful regardless of where the
+  target's scale currently sits. `RegressionHead` is unaffected and still available;
+  pass `MLMWrapper(..., head=RegressionHead(dim))` to keep the previous behaviour.
+
+## [4.3.1] — 2026-07-17
+
+### Fixed
+
+- **`MLMWrapper`'s masked pass leaked gradient to whatever produced `input.x`, despite
+  the target being detached.** `corrupted_x` was built from the live (undetached)
+  `input.x` at unmasked positions; self-attention mixes those into the masked
+  positions' predictions, so `mlm_loss` still reached back through every unmasked
+  position — confirmed empirically at 432/512 nonzero gradient entries under a
+  realistic 15% mask ratio. The existing test only covered 100% masking, which zeroes
+  that path as a side effect of `torch.where`'s backward and never exercised the
+  realistic partial-masking case.
+
+  Fixed by detaching `input.x` once, up front
+  (`input = input._replace(x=input.x.detach())`), so both `corrupted_x` and `target`
+  are built from the same severed copy. `mlm_loss` now trains only `mask_token`,
+  `encoder`, and `head` — never whatever produced `input.x` — regardless of mask
+  ratio. Added a regression test at the realistic mask ratio
+  (`test_wrapper_severs_gradient_to_upstream_input_under_partial_masking`) alongside
+  the existing full-masking one.
+
 ## [4.3.0] — 2026-07-17
 
 ### Added

stackformers-5.0.0rc3/Justfile

@@ -0,0 +1,73 @@
+set dotenv-load := false
+
+python_sources := "stackformers/ tests/ examples/"
+
+# List repository commands
+default:
+    @just --list
+
+# Install all dependencies (including dev)
+[group('environment')]
+sync:
+    uv sync --group dev
+
+# Build distributions without retaining stale artifacts
+[group('packaging')]
+build:
+    uv build --clear
+
+# Format source and tests
+[group('quality')]
+fmt:
+    uv run ruff format {{ python_sources }}
+
+# Check formatting without modifying files
+[group('quality')]
+fmt-check:
+    uv run ruff format --check {{ python_sources }}
+
+# Lint source and tests
+[group('quality')]
+lint:
+    uv run ruff check {{ python_sources }}
+
+# Lint and auto-fix what's safe
+[group('quality')]
+lint-fix:
+    uv run ruff check --fix {{ python_sources }}
+
+# Type-check with pyrefly
+[group('quality')]
+types:
+    uv run pyrefly check {{ python_sources }}
+
+# Run tests
+[group('test')]
+test:
+    uv run pytest tests/ -q
+
+# Run tests with coverage report
+[group('test')]
+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
+[group('test')]
+test-only target:
+    uv run pytest {{ target }} -v
+
+# Run fmt-check + lint + types + test (CI gate)
+[group('quality')]
+check: fmt-check lint types test
+
+# Format, lint-fix, then run all checks
+[group('quality')]
+fix: fmt lint-fix
+    just check
+
+# Remove build artifacts and caches
+[group('maintenance')]
+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-4.3.0 → stackformers-5.0.0rc3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: stackformers
-Version: 4.3.0
+Version: 5.0.0rc3
 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
@@ -43,7 +43,9 @@ Requires-Dist: einops>=0.7
 Requires-Dist: jaxtyping>=0.2
 Requires-Dist: numpy>=2.4.4
 Requires-Dist: pydantic>=2.0
+Requires-Dist: returns<0.29,>=0.28.0
 Requires-Dist: torch>=2.2
+Requires-Dist: typing-extensions>=4.12
 Description-Content-Type: text/markdown
 
 # stackformers
@@ -55,6 +57,10 @@ Every architectural choice — positional encoding, normalization, feedforward v
 uv add stackformers
 ```
 
+Runnable examples cover variable-length layouts, causal language modeling, image patches,
+continuous point coordinates, user subclassing, and cached ONNX Runtime deployment. See the
+[`examples` guide](examples/README.md).
+
 ---
 
 ## Why
@@ -108,6 +114,33 @@ from stackformers import windowed_encoder_config
 windowed_encoder_config(dim=512, heads=8, num_layers=6, window_size=128)
 ```
 
+Variable-width encoder — one model and head dimension per Transformer block:
+
+```python
+from stackformers import VariableWidthTransformerEncoder, variable_width_encoder_config
+
+config = variable_width_encoder_config(
+    d_models=[512, 512, 768, 768, 1024],
+    dim_heads=[64, 64, 64, 64, 128],
+)
+model = VariableWidthTransformerEncoder(config)
+```
+
+Omitting `heads` derives each block's head count as `d_model // dim_head`. Set `heads=12` or
+`heads=[12, 12, 16, 16, 16]` to choose the internal attention width independently of `d_models`.
+`dim_heads`, `heads`, `causal`, `ff_mult`, `dropout`, and `norm_placement` accept either a shared
+value or a list with one entry per block. The preset inserts a bias-free
+learned projection before a block when its residual width differs from the preceding block;
+equal-width neighbors use an identity. The convenience factory expands each width pair into a
+`VariableWidthEncoderLayerConfig` whose attention, feed-forward, norm, positional encoding, and
+attention bias are all explicit. Construct these layer configs directly to replace any component;
+the model does not rebuild or override them. The input width must match the first layer's
+`attn.dim`, and the output uses the last layer's width.
+
+Reference: [Wu et al., “Variable-Width Transformers” (2026)](https://arxiv.org/abs/2606.18246).
+The paper uses parameter-free residual resizing; this preset instead makes each width transition
+a learned linear projection, so it is a related experimental variant rather than an exact replica.
+
 Encoder–decoder:
 
 ```python
@@ -117,6 +150,43 @@ 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))
 ```
 
+For autoregressive export, build cross K/V once and pass both required cache tensors to every
+one-token decoder invocation:
+
+```python
+import torch
+
+from stackformers import (
+    CachedDecoderWrapper,
+    DecoderCrossAttentionCacheBuilder,
+)
+
+model.eval()
+cache_builder = DecoderCrossAttentionCacheBuilder(model).eval()
+cached_decoder = CachedDecoderWrapper(model).eval()
+
+cross_cache = cache_builder(context_input)
+self_kv_cache = torch.empty(num_layers, 2, batch, kv_heads, 0, dim_head)
+step_i = torch.tensor([0], dtype=torch.int64)
+cache_onnx = torch.onnx.export(cache_builder, (context_input,), dynamo=True)
+decoder_onnx = torch.onnx.export(
+    cached_decoder,
+    (target_token, cross_cache, self_kv_cache, step_i),
+    dynamo=True,
+)
+```
+
+The cache builder runs once per encoder context. Each decoder call consumes exactly one target
+token, the immutable cross cache, the growing self cache, and `step_i`; it returns the output token
+and a self cache extended by one position. K/V tensors retain `kv_heads` rather than expanded query
+heads. Cache construction positions cross K once, while cached decoding positions each new self K
+and both self/cross Q, so all built-in positional encodings remain available. Batch, source, and
+past-target axes are dynamic during export; model width, layer count, head geometry, and the
+one-token axis remain static because weights constrain them.
+
+See the complete eager, ONNX export, and raw ONNX Runtime demonstration in
+[`examples/encoder_decoder_onnxruntime.py`](examples/encoder_decoder_onnxruntime.py).
+
 ### Explicit config
 
 Full control with JSON round-trip via `kind` discriminators:
@@ -141,6 +211,24 @@ model = TransformerEncoder(cfg)
 cfg2 = TransformerEncoderConfig.model_validate(cfg.model_dump())
 ```
 
+`norm_placement` is available on `TransformerEncoderConfig`, `TransformerDecoderConfig`, and
+`CrossAttenderConfig`, as well as their plain-config helpers. It accepts four layouts and defaults
+to `"pre"`, so existing constructors, serialized configs, checkpoints, outputs, and gradients retain
+the previous behavior when the field is omitted. For a residual branch `F` with norm `N`:
+
+| Value | Branch equation | Reference |
+| --- | --- | --- |
+| `"pre"` | `x + F(N(x))` | [Xiong et al., 2020](https://proceedings.mlr.press/v119/xiong20b.html) |
+| `"post"` | `N(x + F(x))` | [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762) |
+| `"sandwich"` | `x + N_post(F(N_pre(x)))` | [Ding et al., 2021](https://arxiv.org/abs/2105.13290) |
+| `"reordered"` | `x + N(F(x))` | [Liu et al., 2022](https://arxiv.org/abs/2111.09883); [OLMo Team et al., 2025](https://arxiv.org/abs/2501.00656) |
+
+Sandwich placement creates independent pre- and post-branch norms. Reordered placement follows
+the OLMo 2 residual layout; enable QK-Norm separately in the attention config when reproducing the
+broader OLMo 2 stabilization recipe. Preset construction maps each value to a focused encoder,
+decoder, or cross-attender layer class. Decoder and cross-attender placement applies to the
+target/query residual stream; the context sequence is not normalized or mutated by these layers.
+
 ### Custom wiring
 
 Wire layers yourself when presets aren't enough:
@@ -176,7 +264,7 @@ encoder = Encoder(layers=layers, final_norm=RMSNorm(RMSNormConfig(dim=512)))
 | 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 |
+| Feedforward | SwiGLU, HardSwishGLU, GEGLU, GELU, ReLU² |
 | Normalization | RMSNorm, LayerNorm |
 | Presets | Encoder, Decoder, CrossAttender |
 
@@ -197,6 +285,29 @@ just test     # test
 just check    # full CI gate
 ```
 
+The quality gate checks the library, tests, and examples against Python 3.11. Ruff enforces
+annotations, common bug patterns, typing imports, and qualified, current suppressions.
+Pyrefly uses the strict preset with explicit/implicit `Any`, untyped returns, missing
+annotations, and stale suppressions treated as errors. Tool minimums are recorded in both
+the development dependencies and checker configs.
+
+Module classes expose `__call__ = forward` only under `TYPE_CHECKING`. This gives callers the
+declared input/output contract while leaving PyTorch's hooks and runtime dispatch intact.
+Local `Tensor` annotations at built-in PyTorch projection boundaries rely on those modules'
+tensor-return contracts; they do not validate tensor shapes. Experimental kernel results remain
+`object` until runtime validation. Export shape trees use opaque `object` leaves because
+PyTorch validates their module-dependent structure.
+
+Core admission checks return the shared `returns.result.Result[T, E]` type using `Success`
+and `Failure`, so other libraries can inspect the same outcomes without converting custom
+records. Public adapters preserve tensor returns and raise the stored exception with its
+original type, message, and cause. Pydantic still presents configuration failures as
+`ValidationError`. The experimental backend returns a typed fallback reason internally;
+its existing adapter retains tensor-or-`None` returns and warning behavior. Expected failures
+are explicit in these core signatures; unexpected Python or PyTorch exceptions can still
+propagate. This path uses ordinary constructors, matching, and typed access supported by
+Pyrefly; it requires no mypy plugin.
+
 ---
 
 ## License