diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..5dede5a
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,20 @@
+name: ci
+
+on:
+ push:
+ branches: [main]
+ pull_request:
+
+jobs:
+ test:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+ - uses: actions/setup-python@v5
+ with:
+ python-version: "3.11"
+ cache: pip
+ - run: python -m pip install --upgrade pip
+ - run: python -m pip install -e ".[test]"
+ - run: python -m compileall -q model
+ - run: pytest
diff --git a/README.md b/README.md
index c72ee7b..4f31418 100644
--- a/README.md
+++ b/README.md
@@ -1,431 +1,78 @@
-
+# OUROBOROS
-
-
-# ๐ OUROBOROS
-
-### *Optimal Unified Residual Operations with Bounded Orthogonal Reflection and Spectral Control*
-
-**Teaching Neural Networks the Art of Forgetting**
-
-[](https://www.python.org/downloads/)
-[](https://pytorch.org/)
-[](https://creativecommons.org/licenses/by/4.0/)
-[](https://huggingface.co/)
-
-[**Overview**](#-overview) โข [**Architecture**](#%EF%B8%8F-architecture) โข [**Installation**](#-installation) โข [**Quick Start**](#-quick-start) โข [**Theory**](#-mathematical-foundations) โข [**Insights**](#-key-insights)
-
----
-
-*"The serpent that devours itself to be reborn โ features consumed, transformed, and emerged anew."*
-
-
-
----
-
-## ๐ Overview
-
-
-
-|
-
-### The Problem
-
-Standard residual networks can only **add** information. They lack the ability to **erase**, **forget**, or **reflect** โ leading to **residual accumulation** where noisy features persist indefinitely.
-
-### Our Solution
-
-A **geometric residual connection** that learns:
-- ๐ง **When to remember** โ preserve critical features
-- ๐๏ธ **When to forget** โ erase noise and outdated info
-- ๐ **When to transform** โ flip representations
-
- |
-
-
-```
-โโโโโโโโโโโโโโโโโโโโโโโโโโโ
-โ Traditional ResNet โ
-โ โ
-โ X_{l+1} = X_l + F(X) โ
-โ โ
-โ โ Can only ADD โ
-โโโโโโโโโโโโโโโโโโโโโโโโโโโ
-
-โโโโโโโโโโโโโโโโโโโโโโโโโโโ
-โ OUROBOROS โ
-โ โ
-โ X_{l+1} = AยทX + ฮฒยทkยทvแต โ
-โ โ
-โ โ
ADD, ERASE, REFLECT โ
-โโโโโโโโโโโโโโโโโโโโโโโโโโโ
-```
-
- |
-
-
-
-**OUROBOROS** enables neural networks to:
-
-| Capability | Description |
-|------------|-------------|
-| โจ **Selective Forgetting** | Surgically erase outdated or noisy information |
-| ๐ **Feature Reflection** | Model oscillatory and oppositional dynamics |
-| ๐ฏ **Spectral Control** | Shape layer-wise transitions with precision |
-| โก **Gradient Stability** | Maintain gradient flow with gated identity |
-
----
-
-## ๐๏ธ Architecture
-
-### The Delta Operator
-
-At the heart of OUROBOROS lies the **Delta Operator** โ a generalized Householder transformation:
-
-
-
-### `A(X) = I โ ฮฒ(X) ยท k(X) ยท k(X)แต`
-
-
-
-| Symbol | Name | Description | Range |
-|:------:|------|-------------|-------|
-| **k(X)** | Reflection Direction | Unit vector defining transformation axis | โkโ = 1 |
-| **ฮฒ(X)** | Scalar Gate | Controls transformation intensity | [0, 2] |
-| **v(X)** | Value Vector | New information to inject | โแตแต |
-
----
-
-### โก The Magic of ฮฒ โ One Scalar, Three Transformations
-
-
-
-
-
-
-
-A **single learnable scalar** dynamically interpolates between three geometric operations:
-
-| ฮฒ Value | Transformation | Eigenvalue | Effect |
-|:-------:|----------------|:----------:|--------|
-| **ฮฒ โ 0** | Identity | ฮป = 1 | Pass through unchanged |
-| **ฮฒ โ 1** | Projection | ฮป = 0 | Erase component along k |
-| **ฮฒ โ 2** | Reflection | ฮป = -1 | Flip direction along k |
-
----
-
-### ๐ Geometric Visualization
-
-
-
-
-
-*Vector **v** transformed by Delta Operator. **P(v)** = projection (ฮฒ=1), **R(v)** = reflection (ฮฒ=2). Vector **k** = hyperplane normal.*
-
-
-
----
-
-### ๐ Data Flow
-
-
-
-
-
-
-
-The input **X** splits into three learnable branches that compute **k**, **ฮฒ**, and **v**, which combine through the Delta operation with a skip connection.
-
----
-
-### ๐๏ธ Full Model Architecture
-
-
-
-
-
-
-
-Each **OuroborosBlock** contains:
-- **RMSNorm** โ **Attention** โ **Ouroboros Residual**
-- **RMSNorm** โ **MLP** โ **Ouroboros Residual**
-
----
-
-## ๐งฌ The Ouroboros Residual Block
-
-
-
-
-
-
-
-### Core Update Rule โ The Delta Rule
-
-
-
-```
-X_{l+1} = X_l + ฮฒ ยท k ยท (vแต โ kแต ยท X_l)
- โ โ
- TARGET CURRENT
- (what to write) (what exists)
-```
-
-
-
-This unifies three operations with a **single gate**:
-
-| Operation | Formula | Effect |
-|-----------|---------|--------|
-| **Erasure** | `โฮฒ ยท k ยท (kแต ยท X)` | Removes component along k |
-| **Writing** | `+ฮฒ ยท k ยท vแต` | Injects new information |
-| **Sync** | Same `ฮฒ` | Both scale together |
-
----
-
-## ๐ Mathematical Foundations
-
-### Spectral Decomposition Theorem
-
-> **Theorem**: *For `A = I โ ฮฒยทkยทkแต` where `โkโ = 1`:*
-
-
+OUROBOROS is a **research prototype** for a decoder-only language model that
+replaces additive residual updates with a learned geometric delta update:
```
-ฯ(A) = { 1, 1, ..., 1, (1โฮฒ) }
- โโโโโโฌโโโโโ
- (dโ1) times
+x_next = x + beta * k * (v - )
```
-
-
-| Property | Formula | Notes |
-|----------|---------|-------|
-| **Eigenvalue along k** | `ฮป_k = 1 โ ฮฒ` | Controlled by gate |
-| **Eigenvalues in kโฅ** | `ฮป = 1` | Multiplicity: dโ1 |
-| **Determinant** | `det(A) = 1 โ ฮฒ` | Zero at ฮฒ=1 |
-| **Orthogonality** | `AแตA = I` | When ฮฒ โ {0, 2} |
-| **Involution** | `Aยฒ = I` | When ฮฒ = 2 |
-
-### Why Standard ResNets Are Limited
-
-| Property | ResNet | OUROBOROS |
-|----------|:------:|:---------:|
-| Eigenvalues | โ 1 + ฮต | โ [-1, 1] |
-| Negative ฮป | โ No | โ
Yes |
-| Singular | โ No | โ
Yes (ฮฒ=1) |
-| Data-dependent | โ Fixed | โ
Learnable |
+For a fixed unit direction `k`, the linear component has one eigenvalue
+`1 - beta` along `k` and eigenvalue `1` on its orthogonal complement.
+This identity motivates the design; it is not evidence of better training,
+accuracy, forgetting, or gradient stability.
----
+> [!IMPORTANT]
+> This repository does not ship checkpoints, datasets, training recipes, or
+> reproduced comparisons against standard residual networks. Empirical benefits
+> remain unvalidated until controlled experiments are published.
-## ๐ Feature Coupling
+## What is implemented
-
+- Hugging Face `PreTrainedModel` and `PretrainedConfig` integration
+- causal scaled dot-product attention, RoPE, QK RMSNorm, and SwiGLU
+- bounded per-token beta gates and scalar write targets
+- optional causal one-token key/value shift reference
+- tied token embedding / language-model head
+- import, forward/backward, configuration, and block-size regression tests
-
+## Install and verify
-
-
-**Key Insight**: The geometric coherence term `k_i ยท k_j` enables learned feature interactions without explicit cross-attention.
-
----
-
-## ๐ Installation
+Python 3.10+ is required.
```bash
-# Clone the repository
-git clone https://github.com/DivyamTalwar/OUROBOROS.git
-cd ouroboros
-
-# Install dependencies
-pip install torch>=2.0 transformers einops
-
-# Install in development mode
-pip install -e .
+python -m venv .venv
+source .venv/bin/activate
+pip install -e ".[test]"
+pytest
```
-### Requirements
-
-| Package | Version | Purpose |
-|---------|---------|---------|
-| Python | โฅ 3.8 | Runtime |
-| PyTorch | โฅ 2.0 | Deep learning |
-| Transformers | โฅ 4.30 | HuggingFace |
-| einops | โฅ 0.6 | Tensor ops |
-
----
-
-## โก Quick Start
-
-### Basic Usage
+## Minimal example
```python
-from model.ouroboros import OuroborosModel, OuroborosConfig
import torch
+from model import OuroborosConfig, OuroborosModel
-# Configure the model
config = OuroborosConfig(
- vocab_size=50304,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=6,
- head_dim=128,
+ vocab_size=256,
+ hidden_size=64,
+ num_hidden_layers=2,
+ num_attention_heads=4,
+ head_dim=16,
+ block_size=128,
)
-
-# Initialize model
model = OuroborosModel(config)
-print(f"Parameters: {model.get_num_params():,}")
-
-# Forward pass
-input_ids = torch.randint(0, 50304, (2, 512))
-labels = input_ids.clone()
-
-logits, loss = model(input_ids, targets=labels)
-print(f"Loss: {loss.item():.4f}")
-```
-
-### Training Loop
-
-```python
-from torch.optim import AdamW
-
-model = OuroborosModel(config).cuda()
-optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
-
-for batch in dataloader:
- input_ids, labels = batch['input_ids'].cuda(), batch['labels'].cuda()
-
- logits, loss = model(input_ids, targets=labels)
-
- loss.backward()
- optimizer.step()
- optimizer.zero_grad()
-
- print(f"Loss: {loss.item():.4f}")
-```
-
----
-
-## ๐ง Configuration Reference
-
-### Model Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|:----:|:-------:|-------------|
-| `vocab_size` | int | 50304 | Vocabulary size |
-| `hidden_size` | int | 768 | Model dimension |
-| `num_hidden_layers` | int | 12 | Number of blocks |
-| `num_attention_heads` | int | 6 | Attention heads |
-| `head_dim` | int | 128 | Dimension per head |
-| `block_size` | int | 1024 | Max sequence length |
-
-### Ouroboros Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|:----:|:-------:|-------------|
-| `ouroboros_k_eps` | float | 1e-6 | k normalization ฮต |
-| `ouroboros_beta_init` | float | 1.0 | Initial ฮฒ [0, 2] |
-| `ouroboros_v_sigmoid` | bool | True | Sigmoid on v |
-| `ouroboros_v_sigmoid_scale` | float | 4.0 | v scale factor |
-
----
-
-## ๐ Project Structure
-
-```
-OUROBOROS/
-โโโ ๐ README.md # Documentation
-โโโ assets/ # Images
-โ โโโ banner.png
-โ โโโ architecture.png
-โ โโโ beta_spectrum.png
-โ โโโ geometric_transform.png
-โ โโโ dataflow.png
-โ โโโ model_architecture.png
-โ โโโ feature_coupling.png
-โโโ ๐ model/
- โโโ ouroboros.py # Core implementation
-```
-
----
-
-## ๐ก Key Insights
-
-### Why OUROBOROS Works
-
-| Challenge | ResNet | OUROBOROS |
-|-----------|:------:|:---------:|
-| Noisy features accumulate | โ Can only add | โ
Can erase |
-| Oscillatory patterns | โ No negative ฮป | โ
ฮป โ [-1, 1] |
-| Feature interference | โ No filter | โ
Projection |
-| Gradient stability | โ
Identity | โ
Gated identity |
-
-### Depth-Wise Delta Rule
-
-OUROBOROS is the **depth-wise dual** of time-wise recurrence:
-
-```
-Time (DeltaNet): S_t = A ยท S_{t-1} + ฮฒ ยท k ยท vแต
-Depth (OUROBOROS): X_{l+1} = A ยท X_l + ฮฒ ยท k ยท vแต
-```
-
----
-
-## ๐ฌ Advanced Topics
-
-### Invertibility
-
-When `ฮฒ โ 1`, the Delta Operator is invertible:
-
-```
-Aโปยน = I + (ฮฒ / (1โฮฒ)) ยท k ยท kแต
+tokens = torch.randint(0, config.vocab_size, (2, 32))
+logits, loss = model(tokens, targets=tokens)
+loss.backward()
+print(logits.shape, float(loss))
```
-At `ฮฒ = 2`: **A = Aโปยน** (orthogonal involution).
-
----
-
-## ๐ Citation
-
-```bibtex
-@software{ouroboros2025,
- title = {OUROBOROS: Optimal Unified Residual Operations with
- Bounded Orthogonal Reflection and Spectral Control},
- year = {2025},
- url = {https://github.com/DivyamTalwar/OUROBOROS}
-}
-```
-
----
-
-## ๐ค Contributing
-
-1. **Fork** the repository
-2. **Create** feature branch: `git checkout -b feature/amazing`
-3. **Commit** changes: `git commit -m 'Add feature'`
-4. **Push**: `git push origin feature/amazing`
-5. **Open** a Pull Request
-
----
-
-## ๐ License
-
-**Creative Commons Attribution 4.0 International (CC-BY-4.0)**
-
----
-
-
-
-## ๐ OUROBOROS
-
-*The ancient serpent eating its own tail โ a symbol of cyclical transformation.*
+`hidden_size` must equal `num_attention_heads * head_dim`. Inputs longer than
+`block_size` fail explicitly. `crop_block_size(n)` only decreases that limit.
-*Features are consumed, transformed, and reborn through each layer.*
+## Research status and benchmark policy
----
+The spectral statement above is an exact fixed-`k`, fixed-`beta` property.
+The full model is data-dependent and nonlinear. Any empirical claim should be
+backed by an additive-residual control with matched parameters/FLOPs, identical
+data order, multiple seeds, learning curves, downstream evaluation, and
+activation/gradient diagnostics.
-**Built with ๐ for the ML community**
+The optional `ShiftLinear` is a zero-padded, one-token causal reference. Its
+streaming-cache interface is intentionally unsupported until cache correctness is
+tested.
-[โฌ๏ธ Back to Top](#-ouroboros)
+## License
-
+Apache License 2.0. See [LICENSE](LICENSE).
diff --git a/model/__init__.py b/model/__init__.py
new file mode 100644
index 0000000..4feefcc
--- /dev/null
+++ b/model/__init__.py
@@ -0,0 +1,5 @@
+"""Public OUROBOROS model API."""
+
+from .ouroboros import OuroborosConfig, OuroborosModel, OuroborosResidual
+
+__all__ = ["OuroborosConfig", "OuroborosModel", "OuroborosResidual"]
diff --git a/model/init_utils.py b/model/init_utils.py
new file mode 100644
index 0000000..acd9f5f
--- /dev/null
+++ b/model/init_utils.py
@@ -0,0 +1,33 @@
+import math
+
+import torch
+from torch import nn
+
+
+def init_gpt_weights(model: nn.Module, config) -> None:
+ """Initialize weights once while preserving explicitly initialized beta gates."""
+
+ embedding_std = float(getattr(config, "embedding_init_std", 0.02))
+ hidden_std = float(getattr(config, "hidden_init_std_factor", 0.5)) / math.sqrt(
+ int(config.hidden_size)
+ )
+ output_std = hidden_std / math.sqrt(int(config.num_hidden_layers))
+ initialized = set()
+
+ with torch.no_grad():
+ for name, module in model.named_modules():
+ weight = getattr(module, "weight", None)
+ if isinstance(weight, torch.Tensor) and id(weight) not in initialized:
+ if isinstance(module, nn.Embedding):
+ weight.normal_(mean=0.0, std=embedding_std)
+ elif isinstance(module, nn.Linear):
+ std = output_std if name.endswith("c_proj") else hidden_std
+ weight.normal_(mean=0.0, std=std)
+ elif isinstance(module, nn.LayerNorm):
+ weight.fill_(1.0)
+ initialized.add(id(weight))
+
+ bias = getattr(module, "bias", None)
+ is_beta_gate = name.endswith(".beta") or name.endswith(".beta_out")
+ if isinstance(bias, torch.Tensor) and not is_beta_gate:
+ bias.zero_()
diff --git a/model/kv_shift.py b/model/kv_shift.py
new file mode 100644
index 0000000..742efd8
--- /dev/null
+++ b/model/kv_shift.py
@@ -0,0 +1,25 @@
+import torch
+from torch import nn
+
+
+class ShiftLinear(nn.Linear):
+ """Causal one-token shift followed by a linear projection.
+
+ The optional cache argument is reserved for a future streaming cache. Passing
+ a non-None cache fails explicitly rather than silently producing wrong states.
+ """
+
+ def __init__(self, in_features, out_features, num_heads, bias=False):
+ super().__init__(in_features, out_features, bias=bias)
+ if num_heads <= 0 or out_features % num_heads:
+ raise ValueError("out_features must be divisible by num_heads")
+ self.num_heads = int(num_heads)
+
+ def forward(self, x: torch.Tensor, cache=None) -> torch.Tensor:
+ if cache is not None:
+ raise NotImplementedError("streaming ShiftLinear cache is not implemented")
+ if x.ndim != 3:
+ raise ValueError("ShiftLinear expects (batch, sequence, features)")
+ shifted = torch.zeros_like(x)
+ shifted[:, 1:] = x[:, :-1]
+ return super().forward(shifted)
diff --git a/model/ouroboros.py b/model/ouroboros.py
index c7ecf80..c6ce542 100644
--- a/model/ouroboros.py
+++ b/model/ouroboros.py
@@ -251,6 +251,12 @@ def __init__(self, config):
init_gpt_weights(self, config)
def forward(self, idx, targets=None, return_logits=True, output_all_seq=False):
+ if idx.ndim != 2:
+ raise ValueError("idx must have shape (batch, sequence)")
+ if idx.size(1) > self.config.block_size:
+ raise ValueError(
+ f"sequence length {idx.size(1)} exceeds block_size {self.config.block_size}"
+ )
x = self.transformer.wte(idx) # token embeddings of shape (b, t, hidden_size)
for block in self.transformer.h:
x = block(x)
@@ -281,8 +287,13 @@ def forward(self, idx, targets=None, return_logits=True, output_all_seq=False):
return logits, loss
def crop_block_size(self, block_size):
- """Model surgery to decrease the block size if necessary."""
- pass
+ """Decrease the configured maximum sequence length."""
+ block_size = int(block_size)
+ if block_size <= 0:
+ raise ValueError("block_size must be positive")
+ if block_size > self.config.block_size:
+ raise ValueError("crop_block_size cannot increase the configured block size")
+ self.config.block_size = block_size
def estimate_mfu(self, fwdbwd_per_iter, dt):
"""Estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS."""
diff --git a/model/pydantic_config.py b/model/pydantic_config.py
new file mode 100644
index 0000000..028f316
--- /dev/null
+++ b/model/pydantic_config.py
@@ -0,0 +1,34 @@
+from dataclasses import fields
+
+
+_POSITIVE_FIELDS = (
+ "vocab_size",
+ "num_hidden_layers",
+ "num_attention_heads",
+ "hidden_size",
+ "head_dim",
+ "block_size",
+)
+
+
+def validate_pretrained_config_kwargs(config_cls, kwargs):
+ """Materialize dataclass defaults and validate OUROBOROS dimensions.
+
+ Unknown values are preserved for Hugging Face PretrainedConfig compatibility.
+ """
+
+ supplied = dict(kwargs)
+ values = {
+ field.name: supplied.pop(field.name, getattr(config_cls, field.name))
+ for field in fields(config_cls)
+ }
+ values.update(supplied)
+
+ for name in _POSITIVE_FIELDS:
+ if int(values[name]) <= 0:
+ raise ValueError(f"{name} must be positive")
+ if int(values["hidden_size"]) != int(values["num_attention_heads"]) * int(values["head_dim"]):
+ raise ValueError("hidden_size must equal num_attention_heads * head_dim")
+ if float(values["ouroboros_k_eps"]) <= 0:
+ raise ValueError("ouroboros_k_eps must be positive")
+ return values
diff --git a/model/rmsnorm.py b/model/rmsnorm.py
new file mode 100644
index 0000000..09c365a
--- /dev/null
+++ b/model/rmsnorm.py
@@ -0,0 +1,21 @@
+import torch
+from torch import nn
+
+
+class RMSNorm(nn.Module):
+ """Numerically stable root-mean-square normalization."""
+
+ def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = True):
+ super().__init__()
+ if dim <= 0:
+ raise ValueError("dim must be positive")
+ self.eps = float(eps)
+ if elementwise_affine:
+ self.weight = nn.Parameter(torch.ones(dim))
+ else:
+ self.register_parameter("weight", None)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ scale = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
+ normalized = (x.float() * scale).to(dtype=x.dtype)
+ return normalized if self.weight is None else normalized * self.weight
diff --git a/model/rotary.py b/model/rotary.py
new file mode 100644
index 0000000..175016e
--- /dev/null
+++ b/model/rotary.py
@@ -0,0 +1,34 @@
+import torch
+from torch import nn
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+ first, second = x.chunk(2, dim=-1)
+ return torch.cat((-second, first), dim=-1)
+
+
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+ if x.shape[-1] % 2:
+ raise ValueError("rotary embedding dimension must be even")
+ return x * cos + _rotate_half(x) * sin
+
+
+class Rotary(nn.Module):
+ """Device/dtype-aware RoPE for tensors shaped (batch, time, heads, dim)."""
+
+ def __init__(self, dim: int, base: float = 10000.0, rope_ratio: float = 1.0):
+ super().__init__()
+ if dim <= 0 or dim % 2:
+ raise ValueError("rotary dimension must be a positive even number")
+ if rope_ratio <= 0:
+ raise ValueError("rope_ratio must be positive")
+ inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+ self.register_buffer("inv_freq", inv_freq / float(rope_ratio), persistent=False)
+
+ def forward(self, x: torch.Tensor):
+ positions = torch.arange(x.shape[1], device=x.device, dtype=self.inv_freq.dtype)
+ frequencies = torch.outer(positions, self.inv_freq.to(device=x.device))
+ embedding = torch.cat((frequencies, frequencies), dim=-1)
+ cos = embedding.cos().to(dtype=x.dtype)[None, :, None, :]
+ sin = embedding.sin().to(dtype=x.dtype)[None, :, None, :]
+ return cos, sin
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..2ea4635
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,22 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "ouroboros-residual"
+version = "0.1.0"
+description = "Research prototype for geometric residual updates"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "Apache-2.0"
+dependencies = ["torch>=2.1", "transformers>=4.40,<5"]
+
+[project.optional-dependencies]
+test = ["pytest>=8"]
+
+[tool.setuptools.packages.find]
+include = ["model*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"
diff --git a/tests/test_model.py b/tests/test_model.py
new file mode 100644
index 0000000..a1cc6fc
--- /dev/null
+++ b/tests/test_model.py
@@ -0,0 +1,54 @@
+import pytest
+import torch
+
+from model import OuroborosConfig, OuroborosModel
+from model.kv_shift import ShiftLinear
+
+
+def tiny_config(**overrides):
+ values = dict(
+ vocab_size=32,
+ num_hidden_layers=2,
+ num_attention_heads=2,
+ hidden_size=16,
+ head_dim=8,
+ block_size=16,
+ using_groupnorm=False,
+ )
+ values.update(overrides)
+ return OuroborosConfig(**values)
+
+
+def test_import_and_forward_backward_smoke():
+ torch.manual_seed(0)
+ model = OuroborosModel(tiny_config())
+ tokens = torch.randint(0, 32, (2, 7))
+ logits, loss = model(tokens, targets=tokens)
+
+ assert logits.shape == (2, 7, 32)
+ assert loss is not None and torch.isfinite(loss)
+ loss.backward()
+ assert model.transformer.wte.weight.grad is not None
+
+
+def test_config_rejects_inconsistent_head_dimensions():
+ with pytest.raises(ValueError, match="must equal"):
+ tiny_config(hidden_size=15)
+
+
+def test_block_size_is_enforced_and_can_only_shrink():
+ model = OuroborosModel(tiny_config())
+ model.crop_block_size(4)
+ assert model.config.block_size == 4
+ with pytest.raises(ValueError, match="exceeds"):
+ model(torch.ones(1, 5, dtype=torch.long))
+ with pytest.raises(ValueError, match="cannot increase"):
+ model.crop_block_size(5)
+
+
+def test_shift_linear_is_causal():
+ layer = ShiftLinear(2, 2, num_heads=1, bias=False)
+ with torch.no_grad():
+ layer.weight.copy_(torch.eye(2))
+ x = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
+ torch.testing.assert_close(layer(x), torch.tensor([[[0.0, 0.0], [1.0, 2.0]]]))