torch-weighttracker 0.1.0__tar.gz

torch_weighttracker-0.1.0/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2026 torch-weighttracker contributors
+
+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.
+

torch_weighttracker-0.1.0/MANIFEST.in

@@ -0,0 +1,6 @@
+include LICENSE
+include README.md
+include pyproject.toml
+include torch_weighttracker/py.typed
+recursive-include tests *.py
+prune tests/ignored

torch_weighttracker-0.1.0/PKG-INFO

@@ -0,0 +1,93 @@
+Metadata-Version: 2.4
+Name: torch-weighttracker
+Version: 0.1.0
+Summary: Tools for tracking structured weight sparsity in PyTorch models.
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/dadyownes15/torch-weighttracker
+Project-URL: Issues, https://github.com/dadyownes15/torch-weighttracker/issues
+Project-URL: Repository, https://github.com/dadyownes15/torch-weighttracker
+Keywords: pytorch,torch,weights,tracking,structured-sparsity,pruning
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: torch>=2.0
+Provides-Extra: structured-bops
+Requires-Dist: fvcore; extra == "structured-bops"
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: fvcore; extra == "dev"
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: ruff>=0.6; extra == "dev"
+Requires-Dist: timm>=1.0; extra == "dev"
+Requires-Dist: torchvision>=0.15; extra == "dev"
+Requires-Dist: twine>=5.0; extra == "dev"
+Dynamic: license-file
+
+# torch-weighttracker
+
+PyTorch tools for tracking structured weight sparsity, regularization signals,
+and bit-operation estimates in neural network modules.
+
+The public API is centered on `WeightTracker`:
+
+```python
+import torch
+from torch import nn
+
+from torch_weighttracker import WeightTracker
+
+model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
+tracker = WeightTracker(model, example_inputs=torch.randn(1, 4))
+print(tracker.view_structures())
+```
+
+## Installation
+
+```bash
+python -m pip install torch-weighttracker
+```
+
+Structured BOPs MAC accounting uses `fvcore` for baseline per-module MACs:
+
+```bash
+python -m pip install "torch-weighttracker[structured-bops]"
+```
+
+## Development
+
+```bash
+python -m pip install -e ".[dev]"
+```
+
+Run tests and lint checks:
+
+```bash
+pytest
+ruff check .
+ruff format --check .
+```
+
+## Smoke Test
+
+```bash
+python -c "from torch_weighttracker import WeightTracker; print(WeightTracker)"
+```
+
+## Status
+
+This package is pre-1.0. Public APIs may still change while the tracker,
+calculation, and regularizer surfaces settle.
+
+## License
+
+MIT

torch_weighttracker-0.1.0/README.md

@@ -0,0 +1,58 @@
+# torch-weighttracker
+
+PyTorch tools for tracking structured weight sparsity, regularization signals,
+and bit-operation estimates in neural network modules.
+
+The public API is centered on `WeightTracker`:
+
+```python
+import torch
+from torch import nn
+
+from torch_weighttracker import WeightTracker
+
+model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
+tracker = WeightTracker(model, example_inputs=torch.randn(1, 4))
+print(tracker.view_structures())
+```
+
+## Installation
+
+```bash
+python -m pip install torch-weighttracker
+```
+
+Structured BOPs MAC accounting uses `fvcore` for baseline per-module MACs:
+
+```bash
+python -m pip install "torch-weighttracker[structured-bops]"
+```
+
+## Development
+
+```bash
+python -m pip install -e ".[dev]"
+```
+
+Run tests and lint checks:
+
+```bash
+pytest
+ruff check .
+ruff format --check .
+```
+
+## Smoke Test
+
+```bash
+python -c "from torch_weighttracker import WeightTracker; print(WeightTracker)"
+```
+
+## Status
+
+This package is pre-1.0. Public APIs may still change while the tracker,
+calculation, and regularizer surfaces settle.
+
+## License
+
+MIT

torch_weighttracker-0.1.0/pyproject.toml

@@ -0,0 +1,77 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch-weighttracker"
+version = "0.1.0"
+description = "Tools for tracking structured weight sparsity in PyTorch models."
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+license-files = ["LICENSE"]
+keywords = ["pytorch", "torch", "weights", "tracking", "structured-sparsity", "pruning"]
+classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "numpy>=1.24",
+    "torch>=2.0",
+]
+
+[project.optional-dependencies]
+structured-bops = [
+    "fvcore",
+]
+dev = [
+    "build>=1.2",
+    "fvcore",
+    "pytest>=8.0",
+    "ruff>=0.6",
+    "timm>=1.0",
+    "torchvision>=0.15",
+    "twine>=5.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/dadyownes15/torch-weighttracker"
+Issues = "https://github.com/dadyownes15/torch-weighttracker/issues"
+Repository = "https://github.com/dadyownes15/torch-weighttracker"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["torch_weighttracker*"]
+
+[tool.setuptools.package-data]
+torch_weighttracker = ["py.typed"]
+
+[tool.pytest.ini_options]
+addopts = [
+    "-ra",
+    "--ignore=tests/test_attention_axes_e2e.py",
+    "--ignore=tests/test_attention_axes_unit.py",
+    "--ignore=tests/test_model_structured_sparsity.py",
+    "--ignore=tests/test_sparsity_controller.py",
+    "--ignore=tests/ignored",
+]
+pythonpath = ["."]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B"]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"

torch_weighttracker-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_weighttracker-0.1.0/tests/__init__.py

@@ -0,0 +1 @@
+

torch_weighttracker-0.1.0/tests/dependency_graph.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+
+import pytest
+import torch.nn as nn
+import torch 
+from torch_weighttracker.torch_pruning.dependency import DependencyGraph
+from torch_weighttracker.torch_pruning.pruner.function import LinearPruner, prune_linear_out_channels 
+
+class simpleMLP(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.first_layer =nn.Linear(1,3)
+        self.secound_layer =nn.Linear(3,3)
+        self.third_layer =nn.Linear(3,1) 
+        self.net = nn.Sequential(
+            self.first_layer,
+            self.secound_layer,
+            self.third_layer
+        )
+
+    def forward(self,x):
+        return self.net(x)
+
+
+def test_specs_creation():
+    model = simpleMLP()
+    input_ex = torch.tensor([[1.0]])
+    print(model.forward(input_ex))
+    graph = DependencyGraph().build_dependency(model = model,  example_inputs=input_ex)
+    all_groups = []
+
+    for group in graph.get_all_groups():
+        all_groups.append(group)
+    # iteraer
+    first_group = all_groups[2] 
+    first_item = first_group[0]
+    print("first item: ",first_item)
+    dep = first_item[0]
+    dep2 = first_group[1][0]
+    print(dep.handler)
+    print(vars(dep))
+    print("dep 2")
+    print(vars(dep2))
+    
+    assert dep.handler == prune_linear_out_channels

torch_weighttracker-0.1.0/tests/fixtures_models.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class TinyTransformerClassifier(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.token_embed = nn.Embedding(32, 16)
+        self.position_embed = nn.Embedding(8, 16)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=16,
+            num_heads=4,
+            batch_first=True,
+        )
+        self.norm1 = nn.LayerNorm(16)
+        self.mlp_in = nn.Linear(16, 32)
+        self.activation = nn.GELU()
+        self.mlp_out = nn.Linear(32, 16)
+        self.norm2 = nn.LayerNorm(16)
+        self.head = nn.Linear(16, 5)
+
+    def forward(self, token_ids: Tensor) -> Tensor:
+        positions = torch.arange(token_ids.size(1), device=token_ids.device).unsqueeze(0)
+        x = self.token_embed(token_ids) + self.position_embed(positions)
+
+        attn_out, _ = self.attn(x, x, x, need_weights=False)
+        x = self.norm1(x + attn_out)
+
+        ff = self.mlp_out(self.activation(self.mlp_in(x)))
+        x = self.norm2(x + ff)
+        return self.head(x[:, 0])
+

torch_weighttracker-0.1.0/tests/test_attention_axes_e2e.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+import torch_structure_analyser as tsa
+from torch_structure_analyser.analysis import StructureAxis
+from tests.fixtures_models import TinyTransformerClassifier
+
+
+class AttentionProbeNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mha = nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
+        self.lin = nn.Linear(8, 8, bias=False)
+
+    def forward(self, x):
+        y, _ = self.mha(x, x, x)
+        return self.lin(y)
+
+
+def _make_probe_controller(model: AttentionProbeNet | None = None) -> tsa.SparsityTracker:
+    if model is None:
+        model = AttentionProbeNet()
+    return tsa.SparsityTracker(
+        model,
+        example_inputs=torch.randn(1, 4, 8),
+        root_module_types=[nn.MultiheadAttention, nn.Linear],
+        num_heads={model.mha: 2},
+        prune_num_heads=True,
+        prune_head_dims=True,
+    )
+
+
+def _make_transformer_controller(model: TinyTransformerClassifier | None = None) -> tsa.SparsityTracker:
+    if model is None:
+        model = TinyTransformerClassifier()
+    token_ids = torch.randint(0, 32, (1, 8))
+    return tsa.SparsityTracker(
+        model,
+        example_inputs=token_ids,
+        root_module_types=[nn.MultiheadAttention, nn.Linear],
+        num_heads={model.attn: model.attn.num_heads},
+        prune_num_heads=True,
+        prune_head_dims=True,
+    )
+
+
+def _zero_mha_out_slices(attention: nn.MultiheadAttention, idxs: list[int]) -> None:
+    embed_dim = attention.embed_dim
+    repeated_idxs = idxs + [idx + embed_dim for idx in idxs] + [idx + 2 * embed_dim for idx in idxs]
+    with torch.no_grad():
+        if attention.q_proj_weight is not None:
+            attention.q_proj_weight[idxs, :] = 0
+        if attention.k_proj_weight is not None:
+            attention.k_proj_weight[idxs, :] = 0
+        if attention.v_proj_weight is not None:
+            attention.v_proj_weight[idxs, :] = 0
+        if attention.in_proj_weight is not None:
+            attention.in_proj_weight[repeated_idxs, :] = 0
+            attention.in_proj_weight[:, idxs] = 0
+        if attention.out_proj is not None:
+            attention.out_proj.weight[idxs, :] = 0
+            attention.out_proj.weight[:, idxs] = 0
+
+
+def _zero_member_slices(module, handler, idxs: list[int]) -> None:
+    with torch.no_grad():
+        if handler in [tsa.prune_linear_out_channels]:
+            module.weight[idxs, :] = 0
+            if module.bias is not None:
+                module.bias[idxs] = 0
+        elif handler in [tsa.prune_linear_in_channels]:
+            module.weight[:, idxs] = 0
+        elif handler in [tsa.prune_layernorm_out_channels]:
+            if module.elementwise_affine:
+                module.weight[idxs] = 0
+                if module.bias is not None:
+                    module.bias[idxs] = 0
+        elif handler in [tsa.prune_batchnorm_out_channels, tsa.prune_groupnorm_out_channels, tsa.prune_instancenorm_out_channels]:
+            if getattr(module, "affine", False):
+                module.weight[idxs] = 0
+                if module.bias is not None:
+                    module.bias[idxs] = 0
+        elif handler in [tsa.prune_embedding_out_channels]:
+            module.weight[:, idxs] = 0
+        elif handler in [tsa.prune_multihead_attention_out_channels]:
+            _zero_mha_out_slices(module, idxs)
+        else:
+            raise AssertionError(f"Unhandled test zeroing rule for {handler.__name__}")
+
+
+def _zero_group_unit(controller: tsa.SparsityTracker, group_id: str, root_indices: list[int]) -> None:
+    group_view = next(view for view in controller.iter_groups() if view.group_id == group_id)
+    root_set = set(root_indices)
+    for member in group_view.members:
+        if not member.measurable:
+            continue
+        local_idxs = [
+            local_idx
+            for local_idx, root_idx in zip(member.local_idxs, member.root_idxs)
+            if root_idx in root_set
+        ]
+        if len(local_idxs) == 0:
+            continue
+        _zero_member_slices(member.module, member.handler, local_idxs)
+
+
+def test_attention_head_dim_report_detects_zeroed_shared_dimension():
+    controller = _make_probe_controller()
+    _zero_group_unit(controller, "lin:prune_in_channels:head_dim", [0, 4])
+
+    report = controller.structured_sparsity()
+
+    assert report.by_group["lin:prune_in_channels:head_dim"].zero_prune_units == ((0, 4),)
+    assert report.by_group["lin:prune_in_channels:head"].stats.removed == 0
+
+
+def test_attention_head_report_detects_zeroed_whole_head():
+    controller = _make_probe_controller()
+    _zero_group_unit(controller, "lin:prune_in_channels:head", [0, 1, 2, 3])
+
+    report = controller.structured_sparsity()
+
+    assert report.by_group["lin:prune_in_channels:head"].zero_prune_units == ((0, 1, 2, 3),)
+    assert report.by_group["lin:prune_in_channels:head_dim"].stats.removed == 0
+
+
+def test_prune_zero_structures_updates_logical_num_heads_after_head_prune():
+    model = TinyTransformerClassifier()
+    controller = _make_transformer_controller(model)
+    _zero_group_unit(controller, "attn:prune_out_channels:head", [0, 1, 2, 3])
+
+    result = controller.prune_zero_structures()
+    attention_views = [view for view in controller.iter_groups() if view.axis == StructureAxis.HEAD]
+
+    assert result.pruned_group_ids == ("attn:prune_out_channels:head",)
+    assert controller.config.num_heads[model.attn] == 3
+    assert attention_views[0].size == 3
+
+
+def test_tiny_transformer_exposes_attention_head_and_head_dim_views_end_to_end():
+    controller = _make_transformer_controller()
+
+    attention_views = [view for view in controller.iter_groups() if view.group_id.startswith("attn:prune_out_channels")]
+
+    assert len(attention_views) == 2
+    assert {view.axis for view in attention_views} == {StructureAxis.HEAD, StructureAxis.HEAD_DIM}

torch_weighttracker-0.1.0/tests/test_attention_axes_unit.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+import torch_structure_analyser as tsa
+from torch_structure_analyser.analysis import (
+    StructureAxis,
+    build_atomic_prune_units,
+    build_grouped_prune_units,
+    build_head_prune_units,
+)
+
+
+class AttentionProbeNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mha = nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
+        self.lin = nn.Linear(8, 8, bias=False)
+
+    def forward(self, x):
+        y, _ = self.mha(x, x, x)
+        return self.lin(y)
+
+
+def _make_controller() -> tsa.SparsityTracker:
+    model = AttentionProbeNet()
+    return tsa.SparsityTracker(
+        model,
+        example_inputs=torch.randn(1, 4, 8),
+        root_module_types=[nn.MultiheadAttention, nn.Linear],
+        num_heads={model.mha: 2},
+        prune_num_heads=True,
+        prune_head_dims=True,
+    )
+
+
+def _zero_attention_indices(model: AttentionProbeNet, idxs: list[int]) -> None:
+    embed_dim = model.mha.embed_dim
+    repeated_idxs = idxs + [idx + embed_dim for idx in idxs] + [idx + 2 * embed_dim for idx in idxs]
+    with torch.no_grad():
+        if model.mha.q_proj_weight is not None:
+            model.mha.q_proj_weight[idxs, :] = 0
+        if model.mha.k_proj_weight is not None:
+            model.mha.k_proj_weight[idxs, :] = 0
+        if model.mha.v_proj_weight is not None:
+            model.mha.v_proj_weight[idxs, :] = 0
+        if model.mha.in_proj_weight is not None:
+            model.mha.in_proj_weight[repeated_idxs, :] = 0
+            model.mha.in_proj_weight[:, idxs] = 0
+        if model.mha.out_proj is not None:
+            model.mha.out_proj.weight[idxs, :] = 0
+            model.mha.out_proj.weight[:, idxs] = 0
+        model.lin.weight[:, idxs] = 0
+
+
+def test_build_atomic_prune_units_returns_singletons():
+    units = build_atomic_prune_units(4)
+    assert [unit.root_indices for unit in units] == [(0,), (1,), (2,), (3,)]
+
+
+def test_build_grouped_prune_units_matches_head_dim_pattern():
+    units = build_grouped_prune_units(8, 2)
+    assert [unit.root_indices for unit in units] == [(0, 4), (1, 5), (2, 6), (3, 7)]
+
+
+def test_build_head_prune_units_returns_contiguous_blocks():
+    units = build_head_prune_units(8, 2)
+    assert [unit.root_indices for unit in units] == [(0, 1, 2, 3), (4, 5, 6, 7)]
+
+
+def test_attention_group_views_include_head_and_head_dim_axes():
+    controller = _make_controller()
+    attention_views = [view for view in controller.iter_groups() if view.group_id.startswith("lin:prune_in_channels")]
+
+    assert [view.axis for view in attention_views] == [StructureAxis.HEAD_DIM, StructureAxis.HEAD]
+    assert [unit.root_indices for unit in attention_views[0].prune_units] == [(0, 4), (1, 5), (2, 6), (3, 7)]
+    assert [unit.root_indices for unit in attention_views[1].prune_units] == [(0, 1, 2, 3), (4, 5, 6, 7)]
+
+
+def test_zero_structure_candidates_prioritize_heads_before_head_dims():
+    controller = _make_controller()
+    _zero_attention_indices(controller.model, list(range(8)))
+
+    candidates = controller.zero_structure_candidates()
+    attention_candidates = [candidate for candidate in candidates if candidate.group_id.startswith("lin:prune_in_channels")]
+
+    assert [candidate.axis for candidate in attention_candidates[:2]] == [StructureAxis.HEAD, StructureAxis.HEAD_DIM]
+
+
+def test_group_lasso_axes_filter_keeps_attention_views_separate():
+    controller = _make_controller()
+
+    head_loss, head_named_terms = controller.group_lasso(axes=(StructureAxis.HEAD,))
+    head_dim_loss, head_dim_named_terms = controller.group_lasso(axes=(StructureAxis.HEAD_DIM,))
+
+    assert head_loss.ndim == 0
+    assert head_dim_loss.ndim == 0
+    assert head_loss.requires_grad
+    assert head_dim_loss.requires_grad
+    assert head_named_terms
+    assert head_dim_named_terms
+    assert all("[0]" in key or "[1]" in key for key in head_named_terms)
+    assert len(head_named_terms) < len(head_dim_named_terms)