stackformers 4.3.1__tar.gz → 5.0.0rc8__tar.gz

{stackformers-4.3.1 → stackformers-5.0.0rc8}/.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.0rc8/AGENTS.md

@@ -0,0 +1,144 @@
+# Repository guidance
+
+## What this project is
+
+Typed, composable, SOLID transformer library for PyTorch.
+Behavior comes from injected dependencies, not constructor flags.
+
+## 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.
+
+Always use `uv` for Python commands.
+
+```bash
+just fmt          # ruff format
+just lint         # ruff check
+just types        # pyrefly
+just test         # pytest
+just check        # formatting check, lint, types, and tests
+```
+
+## Design rules
+
+- Encode distinct modes with closed unions, not optional fields in structured records.
+- Use null objects (e.g. `NoPosEncoding`) instead of `if self.x is not None` inside `forward()`.
+- Do not add flags in `__init__` that change `forward()` behavior; introduce a new class.
+- Keep `forward()` traceable by `torch.compile` and `torch.export`; avoid Python control flow
+  on tensors.
+- Keep one concept per file; split files that cannot be described in one sentence.
+- Annotate tensor arguments with jaxtyping shape strings using the dimension vocabulary below.
+- Type collaborators accepted by `__init__` as protocols rather than concrete implementations.
+
+## 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`.
+
+## Expected failures
+
+Use `Result[T, E] = Ok[T] | Err[E]` from `stackformers.result` for recoverable core failures.
+The alternatives are generic NamedTuples; do not add an external Result dependency or use
+unchecked unwrap methods. Match `Ok(value=value)` / `Err(error=error)` with keyword patterns
+for PyTorch tracing compatibility, and use `assert_never` for closed-union exhaustiveness.
+
+Validate admission before construction instead of raising and catching expected failures.
+Keep exception policy at public adapters for operations that cannot proceed, preserving
+existing API contracts. Catch expected vendor exceptions at the vendor boundary, retaining
+their original causes; unexpected programming failures may propagate.
+
+## Dimension vocabulary
+
+| Symbol | Meaning |
+|--------|---------|
+| `b` | Batch size |
+| `n` | Query sequence length |
+| `s` | Key/value sequence length (`n` for self-attention) |
+| `d` | Model dimension |
+| `h` | Number of heads |
+| `dh` | Dimension per head |
+| `w` | Window size |
+| `nt` | Total tokens in a packed sequence |
+| `m` | Number of masked tokens gathered across a batch or pack |
+
+## File organization
+
+- All library code lives directly under `stackformers/`. Breaking changes go in a new
+  top-level package.
+- Keep configs next to the classes they configure, e.g. `attention/config.py` and
+  `feedforward/config.py`.
+- Keep cross-cutting configs (`LayerConfig`, `EncoderConfig`, `DecoderConfig`) in `config.py`.
+- Each module owns a `factory.py` with a `build_*` function that dispatches on config type
+  via `match`.
+- Keep opinionated presets in `presets/`, with one file per preset class.
+- Define protocols in their owning module (`attention/protocols.py`, etc.), not a central file.
+- Implement protocols structurally; implementations do not import the protocols they implement.
+
+## Config and factory conventions
+
+Every union config type uses a `kind: Literal[...]` discriminator field so configs round-trip
+through JSON unambiguously. Define the union as an
+`Annotated[..., Field(discriminator="kind")]` alias in the same `config.py`.
+
+Each module owns its builder in a co-located `factory.py`, with the signature
+`build_*(config: *Config, ...) -> Protocol`. Presets call factories instead of instantiating
+concrete classes directly.
+
+## Module READMEs
+
+Each module subdirectory has a `README.md`. When changing a module's design by adding a
+variant, removing a concept, or changing a dispatch strategy, update its README in the same
+commit. Cover its purpose, key design decisions, and how to extend it with a new variant.
+Do not repeat information obvious from the source, such as file tables or copied signatures.
+
+## Testing
+
+- Mirror source directory structure, with one test file per source file.
+- Use pytest fixtures for module construction, inputs, and sequence objects. Pure transformation
+  and assertion helpers may be ordinary functions.
+- Parametrize compute tests over `device_dtype` (all device/dtype combinations); use `device`
+  alone with float32 for gradient tests.
+- Test shape contracts and numerical behavior where correctness or compatibility requires it.
+- Use `atol(dtype)` from `tests/conftest.py` for numerical tolerances; exact compatibility
+  assertions may require equality instead.
+
+## Dependencies and PyTorch
+
+- Prefer built-in PyTorch implementations over custom ones, especially
+  `F.scaled_dot_product_attention` and `torch.nn.attention.varlen.varlen_attn`.
+- Do not add optional third-party kernel dependencies; kernels use PyTorch.
+- Use Pydantic for configs. Prefer `Field(gt=0)` / `Field(ge=0)` over `@field_validator`
+  for simple bounds.
+- Prefer `nn.Module` activations (`nn.SiLU()`, `nn.GELU()`) over `torch.nn.functional` calls.
+- Before writing PyTorch code, check the installed version with
+  `uv run python -c "import torch; print(torch.__version__)"`.
+- Check `.venv/lib/python3.*/site-packages/torch/nn/modules/` for available implementations.
+
+## x-transformers reference
+
+Use the local `x-transformers/` checkout, when available, for math and implementation details
+only. Do not copy its architecture.
+
+## Versioning
+
+Maintain the version in `pyproject.toml` under `[project].version`, synchronize it with
+`uv.lock`, and expose it at runtime through `stackformers.__version__`.
+
+Use Semantic Versioning:
+
+- **MAJOR** for breaking public API changes.
+- **MINOR** for new backwards-compatible features.
+- **PATCH** for bug fixes or internal changes without API impact.
+
+Bump the version before committing package or behavior changes on `main`. Version bumps on
+feature or work branches require an explicit user instruction. Changes limited to repository
+guidance, tests, or CI do not require a bump.

{stackformers-4.3.1 → stackformers-5.0.0rc8}/CHANGELOG.md

@@ -6,6 +6,281 @@ 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.
 
+## [5.0.0rc8]
+
+### New Features
+
+- Allow decoder context to be omitted or passed as `None`, skipping cross-attention and
+  its normalization across all four layer topologies. Cached decoding accepts a `None`
+  cross cache without requiring context projections. Existing configs and checkpoints
+  remain compatible.
+
+## [5.0.0rc7] — 2026-09-28
+
+### Developers
+
+- Tighten cached attention admission to the existing attention protocols and retain vendor
+  callable parameter/return types during response validation. Give fake kernel requests an
+  explicit TypedDict schema instead of broad `object` annotations.
+
+## [5.0.0rc6] — 2026-09-26
+
+### Backwards Incompatible Changes
+
+#### Core Result values
+
+Core outcomes now use `stackformers.result.Result[T, E]`, the union of generic NamedTuples
+`Ok[T]` and `Err[E]`, instead of `returns.result.Result`. Consumers of core outcomes must
+match `Ok(value=value)` / `Err(error=error)` and access the narrowed payload; `unwrap`,
+`failure`, and `map` methods are no longer available. Public tensor/configuration returns,
+validation exceptions, and backend fallback warnings remain compatible. Remove any reliance
+on Stackformers installing `returns` as a transitive dependency.
+
+### Developers
+
+- Validate cached decoder attention admission before constructing executors, returning
+  expected failures as data instead of raising and catching constructor exceptions.
+
+## [5.0.0rc5] — 2026-09-26
+
+### Developers
+
+- Consolidate repository guidance in `AGENTS.md` and remove `CLAUDE.md`.
+- Reduce the Justfile to eight recipes. Replace `fmt-check`, `lint-fix`, `test-only`, and
+  `test-cov` with arguments to `fmt`, `lint`, and `test`; remove `fix` and `clean`.
+  Keep `check` as the full quality gate and preserve argument boundaries when forwarding flags.
+
+## [5.0.0rc4] — 2026-09-26
+
+### Developers
+
+- Require kernels injected into the experimental attention core to return
+  `Result[Tensor, VarlenFallback]`. Normalize PyTorch responses and expected exceptions
+  in the vendor adapter, preserving public tensor-or-None returns, warnings, and error causes.
+
+## [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

stackformers-5.0.0rc8/Justfile

@@ -0,0 +1,27 @@
+set dotenv-load := false
+set positional-arguments
+
+python_sources := "stackformers/ tests/ examples/"
+
+default:
+    @just --list
+
+sync:
+    uv sync --group dev
+
+build:
+    uv build --clear
+
+fmt *args:
+    uv run ruff format {{ python_sources }} "$@"
+
+lint *args:
+    uv run ruff check {{ python_sources }} "$@"
+
+types:
+    uv run pyrefly check {{ python_sources }}
+
+test *args:
+    uv run pytest -q "$@"
+
+check: (fmt "--check") lint types test