sparse-layers 0.2.0__tar.gz

sparse_layers-0.2.0/.github/workflows/publish.yml

@@ -0,0 +1,26 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Build
+        run: |
+          pip install build
+          python -m build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

sparse_layers-0.2.0/.gitignore

@@ -0,0 +1,11 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+.eggs/
+*.egg
+.pytest_cache/
+.coverage
+htmlcov/
+.pixi/

sparse_layers-0.2.0/CHANGELOG.md

@@ -0,0 +1,17 @@
+# Changelog
+
+## v0.2.0 - 2026-03-23
+
+Initial public release as `sparse-layers` (renamed from `butterfly-layers`).
+
+### Added
+
+- Butterfly-factorized linear layer (`ButterflyLinear`) — O(n log n) parameters as a drop-in replacement for `nn.Linear`
+- `PaddedButterflyLinear` for non-power-of-2 dimensions
+- `ButterflyMLP` and `ButterflyMultiHeadAttention` — sparse MLP and attention building blocks
+- SSE (State Space Exploration) attention modules: `SSEAttention`, `SSEAttentionAdaptive`, `SSEMultiHeadAttention`
+- SSE infrastructure: partition selector, sparse softmax, multi-partition state management
+- `LinearAttention` baseline with O(n) complexity
+- Variable-length sequence operations (`SSEVarlenOps`) and masking utilities (`SSEMaskingOps`)
+- Pydantic-based configuration for all SSE modules
+- Boundary enforcement test — verifies the package has no infrastructure dependencies

sparse_layers-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Philipp Hematty
+
+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.

sparse_layers-0.2.0/PKG-INFO

@@ -0,0 +1,113 @@
+Metadata-Version: 2.4
+Name: sparse-layers
+Version: 0.2.0
+Summary: Structured sparse layers for building memory-efficient neural networks
+Project-URL: Homepage, https://github.com/PhilHem/sparse-layers
+Project-URL: Repository, https://github.com/PhilHem/sparse-layers
+Author-email: Philipp Hematty <34508934+PhilHem@users.noreply.github.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: attention,butterfly,neural-network,pytorch,sparse,sse
+Classifier: Development Status :: 3 - Alpha
+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
+Requires-Python: >=3.11
+Requires-Dist: pydantic<3,>=2.12.4
+Requires-Dist: torch<3,>=2.7.1
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-cov<8,>=7.0.0; extra == 'dev'
+Requires-Dist: pytest-xdist<4,>=3.8.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# sparse-layers
+
+Structured sparse layers for building memory-efficient neural networks on PyTorch. Drop-in replacements for standard layers using butterfly factorization, SSE attention, and other sparse primitives.
+
+## Install
+
+```bash
+pip install sparse-layers
+```
+
+## Usage
+
+```python
+import torch
+from sparse_layers import ButterflyLinear, ButterflyMLP
+
+# Drop-in replacement for nn.Linear with O(n log n) parameters
+layer = ButterflyLinear(in_features=256, out_features=256)
+x = torch.randn(32, 256)
+y = layer(x)
+
+# MLP with butterfly-factorized linear layers
+mlp = ButterflyMLP(in_features=256, hidden_features=512, out_features=256)
+y = mlp(x)
+```
+
+### Attention
+
+```python
+from sparse_layers import ButterflyMultiHeadAttention, MultiHeadAttention
+
+# Standard multi-head attention
+attn = MultiHeadAttention(d_model=256, num_heads=8)
+
+# Butterfly-factorized variant (fewer parameters, same interface)
+bf_attn = ButterflyMultiHeadAttention(d_model=256, num_heads=8)
+
+seq = torch.randn(32, 128, 256)  # (batch, seq_len, d_model)
+out = bf_attn(seq, seq, seq)
+```
+
+### SSE Attention
+
+State Space Exploration modules for efficient sequence modeling with sparse attention patterns.
+
+```python
+from sparse_layers.sse import SSEAttention, SSEAttentionConfig
+
+config = SSEAttentionConfig(d_model=256, num_partitions=4)
+sse = SSEAttention(config)
+
+x = torch.randn(32, 128, 256)
+out = sse(x)
+```
+
+## Modules
+
+### Layers (`sparse_layers.layers`)
+
+| Module | Description |
+|--------|-------------|
+| `ButterflyLinear` | Linear layer using butterfly matrix factorization — O(n log n) parameters instead of O(n²) |
+| `PaddedButterflyLinear` | ButterflyLinear with automatic padding for non-power-of-2 dimensions |
+| `ButterflyMLP` | Two-layer MLP with butterfly-factorized linear layers |
+| `ButterflyMultiHeadAttention` | Multi-head attention with butterfly-factorized Q/K/V projections |
+| `MultiHeadAttention` | Standard multi-head attention (baseline) |
+| `CustomLinear` | Linear layer with pluggable weight initialization |
+| `CustomMLP` | MLP with CustomLinear layers |
+| `SimpleMLP` | Minimal MLP baseline |
+
+### SSE (`sparse_layers.sse`)
+
+| Module | Description |
+|--------|-------------|
+| `SSEAttention` | Sparse attention with state-space-inspired partitioning |
+| `SSEAttentionAdaptive` | SSE with adaptive implementation selection (naive/batched) |
+| `SSEMultiHeadAttention` | Multi-head variant of SSE attention |
+| `SSEMultiPartitionState` | Manages partition states across sequence chunks |
+| `SSEPartitionSelector` | Selects active partitions per query position |
+| `SSESparseSoftmax` | Sparse softmax over selected partitions |
+| `LinearAttention` | Linear attention baseline (O(n) complexity) |
+| `SSEMaskingOps` | Masking utilities for variable-length SSE |
+| `SSEVarlenOps` | Variable-length sequence operations |
+
+## License
+
+MIT

sparse_layers-0.2.0/README.md

@@ -0,0 +1,87 @@
+# sparse-layers
+
+Structured sparse layers for building memory-efficient neural networks on PyTorch. Drop-in replacements for standard layers using butterfly factorization, SSE attention, and other sparse primitives.
+
+## Install
+
+```bash
+pip install sparse-layers
+```
+
+## Usage
+
+```python
+import torch
+from sparse_layers import ButterflyLinear, ButterflyMLP
+
+# Drop-in replacement for nn.Linear with O(n log n) parameters
+layer = ButterflyLinear(in_features=256, out_features=256)
+x = torch.randn(32, 256)
+y = layer(x)
+
+# MLP with butterfly-factorized linear layers
+mlp = ButterflyMLP(in_features=256, hidden_features=512, out_features=256)
+y = mlp(x)
+```
+
+### Attention
+
+```python
+from sparse_layers import ButterflyMultiHeadAttention, MultiHeadAttention
+
+# Standard multi-head attention
+attn = MultiHeadAttention(d_model=256, num_heads=8)
+
+# Butterfly-factorized variant (fewer parameters, same interface)
+bf_attn = ButterflyMultiHeadAttention(d_model=256, num_heads=8)
+
+seq = torch.randn(32, 128, 256)  # (batch, seq_len, d_model)
+out = bf_attn(seq, seq, seq)
+```
+
+### SSE Attention
+
+State Space Exploration modules for efficient sequence modeling with sparse attention patterns.
+
+```python
+from sparse_layers.sse import SSEAttention, SSEAttentionConfig
+
+config = SSEAttentionConfig(d_model=256, num_partitions=4)
+sse = SSEAttention(config)
+
+x = torch.randn(32, 128, 256)
+out = sse(x)
+```
+
+## Modules
+
+### Layers (`sparse_layers.layers`)
+
+| Module | Description |
+|--------|-------------|
+| `ButterflyLinear` | Linear layer using butterfly matrix factorization — O(n log n) parameters instead of O(n²) |
+| `PaddedButterflyLinear` | ButterflyLinear with automatic padding for non-power-of-2 dimensions |
+| `ButterflyMLP` | Two-layer MLP with butterfly-factorized linear layers |
+| `ButterflyMultiHeadAttention` | Multi-head attention with butterfly-factorized Q/K/V projections |
+| `MultiHeadAttention` | Standard multi-head attention (baseline) |
+| `CustomLinear` | Linear layer with pluggable weight initialization |
+| `CustomMLP` | MLP with CustomLinear layers |
+| `SimpleMLP` | Minimal MLP baseline |
+
+### SSE (`sparse_layers.sse`)
+
+| Module | Description |
+|--------|-------------|
+| `SSEAttention` | Sparse attention with state-space-inspired partitioning |
+| `SSEAttentionAdaptive` | SSE with adaptive implementation selection (naive/batched) |
+| `SSEMultiHeadAttention` | Multi-head variant of SSE attention |
+| `SSEMultiPartitionState` | Manages partition states across sequence chunks |
+| `SSEPartitionSelector` | Selects active partitions per query position |
+| `SSESparseSoftmax` | Sparse softmax over selected partitions |
+| `LinearAttention` | Linear attention baseline (O(n) complexity) |
+| `SSEMaskingOps` | Masking utilities for variable-length SSE |
+| `SSEVarlenOps` | Variable-length sequence operations |
+
+## License
+
+MIT

sparse_layers-0.2.0/pyproject.toml

@@ -0,0 +1,46 @@
+[project]
+name = "sparse-layers"
+version = "0.2.0"
+requires-python = ">= 3.11"
+description = "Structured sparse layers for building memory-efficient neural networks"
+license = "MIT"
+authors = [
+    { name = "Philipp Hematty", email = "34508934+PhilHem@users.noreply.github.com" },
+]
+readme = "README.md"
+keywords = ["butterfly", "neural-network", "sparse", "pytorch", "attention", "sse"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "torch>=2.7.1,<3",
+    "pydantic>=2.12.4,<3",
+]
+
+[project.urls]
+Homepage = "https://github.com/PhilHem/sparse-layers"
+Repository = "https://github.com/PhilHem/sparse-layers"
+
+[project.optional-dependencies]
+dev = ["pytest", "pytest-xdist>=3.8.0,<4", "pytest-cov>=7.0.0,<8"]
+
+[build-system]
+build-backend = "hatchling.build"
+requires = ["hatchling"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/sparse_layers"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-v --tb=short"
+
+[tool.coverage.run]
+source = ["src/sparse_layers"]
+omit = ["*/tests/*"]

sparse_layers-0.2.0/src/sparse_layers/__init__.py

@@ -0,0 +1,69 @@
+"""sparse-layers: structured sparse layers for building memory-efficient neural networks."""
+
+from sparse_layers.layers import (
+    ButterflyLinear,
+    ButterflyMLP,
+    ButterflyMultiHeadAttention,
+    CustomLinear,
+    CustomMLP,
+    MultiHeadAttention,
+    PaddedButterflyLinear,
+    SimpleMLP,
+)
+from sparse_layers.sse import (
+    LinearAttention,
+    LinearAttentionConfig,
+    NaiveMultiPartitionState,
+    NaiveSSEAttention,
+    NaiveSSEMultiHeadAttention,
+    SSEAttention,
+    SSEAttentionAdaptive,
+    SSEAttentionAdaptiveConfig,
+    SSEAttentionConfig,
+    SSEMaskingOps,
+    SSEMaskingOpsConfig,
+    SSEMultiHeadAttention,
+    SSEMultiHeadAttentionConfig,
+    SSEMultiPartitionState,
+    SSEMultiPartitionStateConfig,
+    SSEPartitionSelector,
+    SSEPartitionSelectorConfig,
+    SSESparseSoftmax,
+    SSESparseSoftmaxConfig,
+    SSEVarlenOps,
+    SSEVarlenOpsConfig,
+)
+
+__all__ = [
+    # Layers
+    "ButterflyLinear",
+    "ButterflyMLP",
+    "ButterflyMultiHeadAttention",
+    "CustomLinear",
+    "CustomMLP",
+    "MultiHeadAttention",
+    "PaddedButterflyLinear",
+    "SimpleMLP",
+    # SSE
+    "LinearAttention",
+    "LinearAttentionConfig",
+    "NaiveMultiPartitionState",
+    "NaiveSSEAttention",
+    "NaiveSSEMultiHeadAttention",
+    "SSEAttention",
+    "SSEAttentionAdaptive",
+    "SSEAttentionAdaptiveConfig",
+    "SSEAttentionConfig",
+    "SSEMaskingOps",
+    "SSEMaskingOpsConfig",
+    "SSEMultiHeadAttention",
+    "SSEMultiHeadAttentionConfig",
+    "SSEMultiPartitionState",
+    "SSEMultiPartitionStateConfig",
+    "SSEPartitionSelector",
+    "SSEPartitionSelectorConfig",
+    "SSESparseSoftmax",
+    "SSESparseSoftmaxConfig",
+    "SSEVarlenOps",
+    "SSEVarlenOpsConfig",
+]

sparse_layers-0.2.0/src/sparse_layers/layers/__init__.py

@@ -0,0 +1,23 @@
+"""Core neural network layers with butterfly factorization support."""
+
+from sparse_layers.layers.butterfly_linear import ButterflyLinear
+from sparse_layers.layers.butterfly_mlp import ButterflyMLP
+from sparse_layers.layers.butterfly_multi_head_attention import (
+    ButterflyMultiHeadAttention,
+)
+from sparse_layers.layers.custom_linear import CustomLinear
+from sparse_layers.layers.custom_mlp import CustomMLP
+from sparse_layers.layers.multi_head_attention import MultiHeadAttention
+from sparse_layers.layers.padded_butterfly_linear import PaddedButterflyLinear
+from sparse_layers.layers.simple_mlp import SimpleMLP
+
+__all__ = [
+    "ButterflyLinear",
+    "ButterflyMLP",
+    "ButterflyMultiHeadAttention",
+    "CustomLinear",
+    "CustomMLP",
+    "MultiHeadAttention",
+    "PaddedButterflyLinear",
+    "SimpleMLP",
+]

sparse_layers-0.2.0/src/sparse_layers/layers/butterfly_linear.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+def _is_power_of_two(value: int) -> bool:
+    return value > 0 and (value & (value - 1) == 0)
+
+
+class ButterflyLinear(nn.Module):
+    """Butterfly factorization based linear layer.
+
+    This layer follows the structure described in section 2.3.1 of the
+    `reducing_memory_requirements_ipu_butterfly.md` reference. It consists of
+    log2(N) stages of block-diagonal 2x2 butterfly factors, each stored as a
+    learnable parameter. The layer currently supports square power-of-two
+    dimensions and acts as a drop-in replacement for :class:`torch.nn.Linear`.
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super().__init__()
+
+        if in_features <= 0 or out_features <= 0:
+            raise ValueError("in_features and out_features must be positive integers")
+
+        if in_features != out_features:
+            raise ValueError("ButterflyLinear requires in_features == out_features")
+
+        if not _is_power_of_two(in_features):
+            raise ValueError(
+                "ButterflyLinear requires dimensions that are a power of two"
+            )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self._depth = int(math.log2(in_features))
+
+        num_blocks = in_features // 2
+        default_dtype = torch.get_default_dtype()
+        identity_block = torch.eye(2, dtype=default_dtype).unsqueeze(0)
+
+        self.factors = nn.ParameterList()
+        for _ in range(self._depth):
+            factor = identity_block.repeat(num_blocks, 1, 1)
+            factor += 0.01 * torch.randn_like(factor)
+            self.factors.append(nn.Parameter(factor))
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features, dtype=default_dtype))
+        else:
+            self.register_parameter("bias", None)
+
+    @classmethod
+    def from_linear(
+        cls,
+        layer: nn.Linear,
+        *,
+        optimization_steps: int = 4000,
+        learning_rate: float = 0.1,
+        tolerance: float = 1e-7,
+        seed: int | None = None,
+    ) -> "ButterflyLinear":
+        """Construct a butterfly layer approximating a dense :class:`nn.Linear`.
+
+        The method uses gradient-based fitting over the canonical basis to match
+        the transformation represented by ``layer``. The bias term is copied
+        directly before optimisation to accelerate convergence.
+        """
+
+        if layer.in_features != layer.out_features:
+            raise ValueError("ButterflyLinear.from_linear requires a square nn.Linear")
+
+        if not _is_power_of_two(layer.in_features):
+            raise ValueError(
+                "ButterflyLinear.from_linear requires dimensions that are a power of two"
+            )
+
+        device = layer.weight.device
+        dtype = layer.weight.dtype
+
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        result = cls(layer.in_features, layer.out_features, bias=layer.bias is not None)
+        result = result.to(device=device, dtype=dtype)
+
+        if result.bias is not None and layer.bias is not None:
+            with torch.no_grad():
+                result.bias.copy_(layer.bias)
+
+        params = list(result.factors.parameters())
+        if not params:
+            return result
+
+        optimizer_adam = torch.optim.Adam(params, lr=learning_rate)
+
+        eye_input = torch.eye(layer.in_features, device=device, dtype=dtype)
+        with torch.no_grad():
+            target = layer(eye_input)
+
+        best_loss = float("inf")
+
+        for step in range(optimization_steps):
+            optimizer_adam.zero_grad()
+            output = result(eye_input)
+            loss = F.mse_loss(output, target)
+            loss.backward()
+            optimizer_adam.step()
+
+            current_loss = loss.item()
+            if current_loss < best_loss:
+                best_loss = current_loss
+
+            if best_loss <= tolerance:
+                break
+
+        if best_loss > tolerance:
+            optimizer_lbfgs = torch.optim.LBFGS(
+                params,
+                lr=1.0,
+                max_iter=100,
+                tolerance_grad=1e-12,
+                tolerance_change=1e-12,
+                line_search_fn="strong_wolfe",
+            )
+
+            def closure() -> Tensor:
+                optimizer_lbfgs.zero_grad()
+                output_lbfgs = result(eye_input)
+                lbfgs_loss = F.mse_loss(output_lbfgs, target)
+                lbfgs_loss.backward()
+                return lbfgs_loss
+
+            for _ in range(20):
+                loss = optimizer_lbfgs.step(closure)
+                best_loss = min(best_loss, loss.item())
+                if best_loss <= tolerance:
+                    break
+
+        return result
+
+    def forward(self, input: Tensor) -> Tensor:
+        if input.shape[-1] != self.in_features:
+            raise ValueError(
+                f"Expected last dimension {self.in_features}, got {input.shape[-1]}"
+            )
+
+        original_shape = input.shape[:-1]
+        x = input.reshape(-1, self.in_features)
+
+        for stage_index, factor in enumerate(self.factors):
+            x = self._apply_stage(x, factor, stage_index)
+
+        if self.bias is not None:
+            x = x + self.bias
+
+        return x.reshape(*original_shape, self.out_features)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"bias={self.bias is not None}"
+        )
+
+    def _apply_stage(self, x: Tensor, factor: Tensor, stage: int) -> Tensor:
+        batch = x.shape[0]
+        n = self.in_features
+        block = 1 << (stage + 1)
+        half = block >> 1
+
+        staged = x.reshape(batch, -1, block)
+        staged = staged.reshape(batch, -1, half, 2)
+        staged = staged.permute(0, 1, 3, 2).contiguous()
+        pairs = staged.reshape(batch, -1, 2)
+
+        transformed = torch.einsum("bnc,ncd->bnd", pairs, factor)
+
+        transformed = transformed.reshape(batch, -1, 2, half)
+        transformed = transformed.permute(0, 1, 3, 2).contiguous()
+        transformed = transformed.reshape(batch, -1, block)
+        return transformed.reshape(batch, n)
+
+    def to_linear(self) -> nn.Linear:
+        """Return a dense :class:`nn.Linear` with identical behaviour."""
+
+        factor_tensor = self.factors[0] if len(self.factors) > 0 else None
+        if factor_tensor is not None:
+            device = factor_tensor.device
+            dtype = factor_tensor.dtype
+        elif self.bias is not None:
+            device = self.bias.device
+            dtype = self.bias.dtype
+        else:
+            device = torch.device("cpu")
+            dtype = torch.get_default_dtype()
+
+        linear = nn.Linear(
+            self.in_features,
+            self.out_features,
+            bias=self.bias is not None,
+            device=device,
+            dtype=dtype,
+        )
+
+        with torch.no_grad():
+            if self.bias is not None:
+                bias_backup = self.bias.data.clone()
+                self.bias.zero_()
+            else:
+                bias_backup = None
+
+            identity = torch.eye(self.in_features, device=device, dtype=dtype)
+            weight_matrix = self(identity)
+
+            if bias_backup is not None:
+                self.bias.copy_(bias_backup)
+                linear.bias.copy_(bias_backup)
+
+            linear.weight.copy_(weight_matrix.t())
+
+        return linear