rasnatune 0.0.1__tar.gz

rasnatune-0.0.1/LICENSE

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

rasnatune-0.0.1/PKG-INFO

@@ -0,0 +1,79 @@
+Metadata-Version: 2.4
+Name: rasnatune
+Version: 0.0.1
+Summary: Runtime compression tuning utilities for PyTorch
+Author: rasnatune contributors
+License-Expression: MIT
+Keywords: pytorch,compression,sparsification,quantization
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Dynamic: license-file
+
+# rasnatune
+
+`rasnatune` is a lightweight PyTorch compression helper based on forward hooks.
+
+## Installation
+
+```bash
+pip install rasnatune
+```
+
+## What It Does
+
+- Applies compression at runtime using hooks (no permanent weight rewrite).
+- Supports quantization and unstructured sparsification.
+- Targets `torch.nn.Conv2d` and `torch.nn.Linear`.
+
+## Quick Start
+
+```python
+import torch
+from rasnatune import Compression, SparseWeightUnstructured
+
+model = torch.nn.Sequential(
+    torch.nn.Linear(128, 64),
+    torch.nn.ReLU(),
+    torch.nn.Linear(64, 10),
+)
+
+wrapped = Compression(model)
+wrapped.attach(
+    SparseWeightUnstructured,
+    filter=lambda m: isinstance(m, torch.nn.Linear),
+    sparsity=0.5,
+)
+
+x = torch.randn(128)
+y = wrapped(x)
+```
+
+## Public API
+
+Top-level exports:
+
+- `rasnatune.Compression`
+- `rasnatune.Compressor`
+- `rasnatune.QuantizeWeight`
+- `rasnatune.QuantizeActivation`
+- `rasnatune.SparseWeightUnstructured`
+- `rasnatune.SparseActivationUnstructured`
+
+## Compression Classes
+
+- `QuantizeWeight(min=-128, max=127)`:
+  Quantizes layer weights only during forward, then restores original weights.
+- `QuantizeActivation(min=-128, max=127)`:
+  Quantizes the first input activation of the layer before forward.
+- `SparseWeightUnstructured(sparsity=0.5)`:
+  Applies unstructured sparsity to layer weights only during forward, then restores.
+- `SparseActivationUnstructured(sparsity=0.5)`:
+  Applies unstructured sparsity to the first input activation before forward.
+
+## Notes
+
+- `Compression.attach` uses `filter=` to choose target modules.
+- If `filter` is omitted, it will try all submodules, and unsupported modules will raise an assertion.
+- Supported module types for built-in compressors are `torch.nn.Conv2d` and `torch.nn.Linear`.

rasnatune-0.0.1/README.md

@@ -0,0 +1,66 @@
+# rasnatune
+
+`rasnatune` is a lightweight PyTorch compression helper based on forward hooks.
+
+## Installation
+
+```bash
+pip install rasnatune
+```
+
+## What It Does
+
+- Applies compression at runtime using hooks (no permanent weight rewrite).
+- Supports quantization and unstructured sparsification.
+- Targets `torch.nn.Conv2d` and `torch.nn.Linear`.
+
+## Quick Start
+
+```python
+import torch
+from rasnatune import Compression, SparseWeightUnstructured
+
+model = torch.nn.Sequential(
+    torch.nn.Linear(128, 64),
+    torch.nn.ReLU(),
+    torch.nn.Linear(64, 10),
+)
+
+wrapped = Compression(model)
+wrapped.attach(
+    SparseWeightUnstructured,
+    filter=lambda m: isinstance(m, torch.nn.Linear),
+    sparsity=0.5,
+)
+
+x = torch.randn(128)
+y = wrapped(x)
+```
+
+## Public API
+
+Top-level exports:
+
+- `rasnatune.Compression`
+- `rasnatune.Compressor`
+- `rasnatune.QuantizeWeight`
+- `rasnatune.QuantizeActivation`
+- `rasnatune.SparseWeightUnstructured`
+- `rasnatune.SparseActivationUnstructured`
+
+## Compression Classes
+
+- `QuantizeWeight(min=-128, max=127)`:
+  Quantizes layer weights only during forward, then restores original weights.
+- `QuantizeActivation(min=-128, max=127)`:
+  Quantizes the first input activation of the layer before forward.
+- `SparseWeightUnstructured(sparsity=0.5)`:
+  Applies unstructured sparsity to layer weights only during forward, then restores.
+- `SparseActivationUnstructured(sparsity=0.5)`:
+  Applies unstructured sparsity to the first input activation before forward.
+
+## Notes
+
+- `Compression.attach` uses `filter=` to choose target modules.
+- If `filter` is omitted, it will try all submodules, and unsupported modules will raise an assertion.
+- Supported module types for built-in compressors are `torch.nn.Conv2d` and `torch.nn.Linear`.

rasnatune-0.0.1/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "rasnatune"
+version = "0.0.1"
+description = "Runtime compression tuning utilities for PyTorch"
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [{ name = "rasnatune contributors" }]
+license = "MIT"
+keywords = ["pytorch", "compression", "sparsification", "quantization"]
+dependencies = ["torch>=2.0"]
+
+[tool.setuptools.packages.find]
+include = ["rasnatune*"]
+exclude = ["__pycache__*"]

rasnatune-0.0.1/rasnatune/__init__.py

@@ -0,0 +1,6 @@
+from .base import Compressor
+from .compression import Compression
+from .quantization import QuantizeActivation, QuantizeWeight
+from .sparsification import SparseActivationUnstructured, SparseWeightUnstructured
+
+__version__ = "0.0.1"

rasnatune-0.0.1/rasnatune/base.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+import torch
+
+
+class Compressor:
+    _backup: torch.Tensor
+    _hooks: list[torch.utils.hooks.RemovableHandle]
+
+    def attach(self, module: torch.nn.Module) -> None:
+        raise NotImplementedError
+
+    def detach(self, module: torch.nn.Module) -> None:
+        for hook in self._hooks:
+            hook.remove()

rasnatune-0.0.1/rasnatune/compression.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+import torch
+
+from .base import Compressor
+
+
+class Compression(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self.model = model
+        self.registrations: list[Compressor] = []
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def attach(self, compressor: type[Compressor], filter: Callable, *args, **kwargs,):
+        for name, module in self.model.named_modules():
+            if not filter(module):
+                continue
+            instance = compressor(*args, **kwargs)
+            instance.attach(module)
+            self.registrations.append(instance)
+
+    def detach(self, compressor: Compressor):
+        compressor.detach(self.model)
+        if compressor in self.registrations:
+            self.registrations.remove(compressor)
+
+    def clear(self):
+        for compressor in list(self.registrations):
+            self.detach(compressor)

rasnatune-0.0.1/rasnatune/quantization.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+
+from .base import Compressor
+
+
+def quantize(x: torch.Tensor, qmin: int, qmax: int) -> torch.Tensor:
+    max_abs = x.detach().abs().max()
+    if max_abs.item() == 0:
+        return x
+    scale = max_abs / max(qmax, -qmin)
+    q = torch.clamp(torch.round(x / scale), qmin, qmax)
+    quantized = q * scale
+    error = x.detach() - quantized
+    return x - error
+
+
+@dataclass
+class QuantizeWeight(Compressor):
+    min: int = -128
+    max: int = 127
+
+    def attach(self, module: torch.nn.Module) -> None:
+        def pre_hook(mod: torch.nn.Module, inputs) -> None:
+            self._backup = mod.weight.data.detach().clone()
+            mod.weight.data.copy_(quantize(mod.weight.data, qmin=self.min, qmax=self.max))
+
+        def post_hook(mod: torch.nn.Module, inputs, output):
+            mod.weight.data.copy_(self._backup)
+            return output
+
+        self._hooks = [
+            module.register_forward_pre_hook(pre_hook),
+            module.register_forward_hook(post_hook)
+        ]
+
+
+@dataclass
+class QuantizeActivation(Compressor):
+    min: int = -128
+    max: int = 127
+
+    def attach(self, module: torch.nn.Module) -> None:
+        def pre_hook(_module: torch.nn.Module, inputs):
+            return (quantize(inputs[0], qmin=self.min, qmax=self.max), *inputs[1:])
+
+        self._hooks = [
+            module.register_forward_pre_hook(pre_hook)
+        ]

rasnatune-0.0.1/rasnatune/sparsification.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from .base import Compressor
+
+
+def sparsify(x: torch.Tensor, sparsity: float) -> torch.Tensor:
+    if not 0.0 <= sparsity <= 1.0:
+        raise ValueError(f"sparsity must be between 0 and 1, got {sparsity}")
+
+    flat_abs = x.detach().abs().flatten()
+    if flat_abs.numel() == 0 or sparsity == 0.0:
+        return x
+    if sparsity == 1.0:
+        return torch.zeros_like(x)
+
+    num_zero = int(flat_abs.numel() * sparsity)
+    if num_zero == 0:
+        return x
+
+    zero_indices = flat_abs.topk(num_zero, largest=False).indices
+    sparsity_mask = torch.ones_like(flat_abs, dtype=torch.bool)
+    sparsity_mask[zero_indices] = False
+    return x * sparsity_mask.view_as(x).to(x.dtype)
+
+
+@dataclass
+class SparseWeightUnstructured(Compressor):
+    sparsity: float = 0.5
+
+    def attach(self, module: torch.nn.Module) -> None:
+        def pre_hook(mod: torch.nn.Module, _inputs) -> None:
+            self._backup = mod.weight.data.detach().clone()
+            mod.weight.data.copy_(
+                sparsify(mod.weight.data, sparsity=self.sparsity)
+            )
+
+        def post_hook(mod: torch.nn.Module, inputs, output):
+            mod.weight.data.copy_(self._backup)
+            return output
+
+        self._hooks = [
+            module.register_forward_pre_hook(pre_hook),
+            module.register_forward_hook(post_hook),
+        ]
+
+
+@dataclass
+class SparseActivationUnstructured(Compressor):
+    sparsity: float = 0.5
+
+    def attach(self, module: torch.nn.Module) -> None:
+        def pre_hook(_module: torch.nn.Module, inputs):
+            return (sparsify(inputs[0], sparsity=self.sparsity), *inputs[1:])
+
+        self._hooks = [
+            module.register_forward_pre_hook(pre_hook),
+        ]

rasnatune-0.0.1/rasnatune.egg-info/PKG-INFO

@@ -0,0 +1,79 @@
+Metadata-Version: 2.4
+Name: rasnatune
+Version: 0.0.1
+Summary: Runtime compression tuning utilities for PyTorch
+Author: rasnatune contributors
+License-Expression: MIT
+Keywords: pytorch,compression,sparsification,quantization
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Dynamic: license-file
+
+# rasnatune
+
+`rasnatune` is a lightweight PyTorch compression helper based on forward hooks.
+
+## Installation
+
+```bash
+pip install rasnatune
+```
+
+## What It Does
+
+- Applies compression at runtime using hooks (no permanent weight rewrite).
+- Supports quantization and unstructured sparsification.
+- Targets `torch.nn.Conv2d` and `torch.nn.Linear`.
+
+## Quick Start
+
+```python
+import torch
+from rasnatune import Compression, SparseWeightUnstructured
+
+model = torch.nn.Sequential(
+    torch.nn.Linear(128, 64),
+    torch.nn.ReLU(),
+    torch.nn.Linear(64, 10),
+)
+
+wrapped = Compression(model)
+wrapped.attach(
+    SparseWeightUnstructured,
+    filter=lambda m: isinstance(m, torch.nn.Linear),
+    sparsity=0.5,
+)
+
+x = torch.randn(128)
+y = wrapped(x)
+```
+
+## Public API
+
+Top-level exports:
+
+- `rasnatune.Compression`
+- `rasnatune.Compressor`
+- `rasnatune.QuantizeWeight`
+- `rasnatune.QuantizeActivation`
+- `rasnatune.SparseWeightUnstructured`
+- `rasnatune.SparseActivationUnstructured`
+
+## Compression Classes
+
+- `QuantizeWeight(min=-128, max=127)`:
+  Quantizes layer weights only during forward, then restores original weights.
+- `QuantizeActivation(min=-128, max=127)`:
+  Quantizes the first input activation of the layer before forward.
+- `SparseWeightUnstructured(sparsity=0.5)`:
+  Applies unstructured sparsity to layer weights only during forward, then restores.
+- `SparseActivationUnstructured(sparsity=0.5)`:
+  Applies unstructured sparsity to the first input activation before forward.
+
+## Notes
+
+- `Compression.attach` uses `filter=` to choose target modules.
+- If `filter` is omitted, it will try all submodules, and unsupported modules will raise an assertion.
+- Supported module types for built-in compressors are `torch.nn.Conv2d` and `torch.nn.Linear`.

rasnatune-0.0.1/rasnatune.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+rasnatune/__init__.py
+rasnatune/base.py
+rasnatune/compression.py
+rasnatune/quantization.py
+rasnatune/sparsification.py
+rasnatune.egg-info/PKG-INFO
+rasnatune.egg-info/SOURCES.txt
+rasnatune.egg-info/dependency_links.txt
+rasnatune.egg-info/requires.txt
+rasnatune.egg-info/top_level.txt
+tests/test_compression.py
+tests/test_quantization.py
+tests/test_sparsification.py

rasnatune-0.0.1/rasnatune.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

rasnatune-0.0.1/rasnatune.egg-info/requires.txt

@@ -0,0 +1 @@
+torch>=2.0

rasnatune-0.0.1/rasnatune.egg-info/top_level.txt

@@ -0,0 +1 @@
+rasnatune

rasnatune-0.0.1/setup.cfg

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

rasnatune-0.0.1/tests/test_compression.py

@@ -0,0 +1,172 @@
+import torch
+import torch.nn.functional as F
+
+try:
+    _TORCHVISION_NMS_LIB = torch.library.Library("torchvision", "DEF")
+    _TORCHVISION_NMS_LIB.define("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor")
+except RuntimeError:
+    _TORCHVISION_NMS_LIB = None
+
+from torchvision.models import efficientnet_b0, resnet50
+
+from rasnatune.compression import Compression
+from rasnatune.quantization import QuantizeActivation, QuantizeWeight
+from rasnatune.sparsification import SparseActivationUnstructured, SparseWeightUnstructured
+
+
+def _supported_module_filter(module: torch.nn.Module) -> bool:
+    return isinstance(module, (torch.nn.Conv2d, torch.nn.Linear))
+
+
+def _supported_modules(model: torch.nn.Module) -> list[torch.nn.Module]:
+    return [module for module in model.modules() if _supported_module_filter(module)]
+
+
+def _make_toy_model() -> torch.nn.Sequential:
+    model = torch.nn.Sequential(
+        torch.nn.Linear(4, 4, bias=False),
+        torch.nn.ReLU(),
+        torch.nn.Linear(4, 2, bias=False),
+    )
+    with torch.no_grad():
+        model[0].weight.copy_(
+            torch.tensor(
+                [
+                    [1.75, -0.25, 0.5, -1.1],
+                    [0.2, -0.6, 1.4, -0.8],
+                    [0.9, 0.3, -0.4, 1.1],
+                    [-1.2, 0.7, 0.6, -0.5],
+                ],
+                dtype=torch.float32,
+            )
+        )
+        model[2].weight.copy_(
+            torch.tensor(
+                [
+                    [0.8, -1.4, 0.6, -0.2],
+                    [-0.5, 0.9, -1.1, 1.3],
+                ],
+                dtype=torch.float32,
+            )
+        )
+    return model
+
+
+def _attach_quantization(wrapped: Compression) -> None:
+    wrapped.attach(QuantizeWeight, filter=_supported_module_filter, min=-8, max=7)
+    wrapped.attach(QuantizeActivation, filter=_supported_module_filter, min=-8, max=7)
+
+
+def _attach_sparsification(wrapped: Compression) -> None:
+    wrapped.attach(SparseWeightUnstructured, filter=_supported_module_filter, sparsity=0.25)
+    wrapped.attach(SparseActivationUnstructured, filter=_supported_module_filter, sparsity=0.25)
+
+
+def _run_compressed_training_step(model: torch.nn.Module, attach_fn) -> None:
+    wrapped = Compression(model)
+    modules = _supported_modules(wrapped.model)
+    first_module = modules[0]
+    original_weight = first_module.weight.detach().clone()
+
+    attach_fn(wrapped)
+
+    assert len(wrapped.registrations) == 2 * len(modules)
+
+    optimizer = torch.optim.SGD(wrapped.parameters(), lr=1e-2)
+    inputs = torch.randn(2, 3, 64, 64, dtype=torch.float32)
+    targets = torch.tensor([1, 3], dtype=torch.long)
+
+    optimizer.zero_grad(set_to_none=True)
+    logits = wrapped(inputs)
+    loss = F.cross_entropy(logits, targets)
+    loss.backward()
+
+    gradients = [parameter.grad for parameter in wrapped.parameters() if parameter.grad is not None]
+    assert gradients
+    assert all(torch.isfinite(gradient).all().item() for gradient in gradients)
+    torch.testing.assert_close(first_module.weight.detach(), original_weight)
+
+    parameters_before_step = [parameter.detach().clone() for parameter in wrapped.parameters()]
+    optimizer.step()
+
+    assert any(
+        not torch.equal(parameter.detach(), before_step)
+        for parameter, before_step in zip(wrapped.parameters(), parameters_before_step)
+    )
+
+    wrapped.clear()
+    assert wrapped.registrations == []
+
+    cleared_logits = wrapped(inputs)
+    assert cleared_logits.shape == logits.shape
+
+
+def test_compression_attach_and_detach_manage_registrations() -> None:
+    model = _make_toy_model()
+    wrapped = Compression(model)
+    x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.float32)
+
+    wrapped.attach(QuantizeWeight, filter=lambda module: isinstance(module, torch.nn.Linear), min=-1, max=1)
+
+    assert len(wrapped.registrations) == 2
+
+    compressed_output = wrapped(x)
+    wrapped.detach(wrapped.registrations[0])
+
+    assert len(wrapped.registrations) == 1
+
+    wrapped.clear()
+    assert wrapped.registrations == []
+
+    restored_output = wrapped(x)
+    reference_output = model(x)
+
+    torch.testing.assert_close(restored_output, reference_output)
+    assert not torch.equal(compressed_output, reference_output)
+
+
+def test_compression_clear_removes_sparse_registrations() -> None:
+    model = _make_toy_model()
+    wrapped = Compression(model)
+
+    _attach_sparsification(wrapped)
+
+    assert len(wrapped.registrations) == 4
+
+    wrapped.clear()
+
+    assert wrapped.registrations == []
+    output = wrapped(torch.randn(1, 4, dtype=torch.float32))
+    assert output.shape == (1, 2)
+
+
+def test_resnet50_quantized_training_step_runs() -> None:
+    torch.manual_seed(0)
+    model = resnet50(weights=None, num_classes=5)
+    model.train()
+
+    _run_compressed_training_step(model, _attach_quantization)
+
+
+def test_efficientnet_b0_quantized_training_step_runs() -> None:
+    torch.manual_seed(0)
+    model = efficientnet_b0(weights=None, num_classes=5)
+    model.train()
+
+    _run_compressed_training_step(model, _attach_quantization)
+
+
+def test_resnet50_sparse_training_step_runs() -> None:
+    torch.manual_seed(0)
+    model = resnet50(weights=None, num_classes=5)
+    model.train()
+
+    _run_compressed_training_step(model, _attach_sparsification)
+
+
+def test_efficientnet_b0_sparse_training_step_runs() -> None:
+    torch.manual_seed(0)
+    model = efficientnet_b0(weights=None, num_classes=5)
+    model.train()
+
+    _run_compressed_training_step(model, _attach_sparsification)

rasnatune-0.0.1/tests/test_quantization.py

@@ -0,0 +1,127 @@
+import torch
+import torch.nn.functional as F
+
+from rasnatune.quantization import QuantizeActivation, QuantizeWeight, quantize
+
+
+def _make_linear() -> torch.nn.Linear:
+    module = torch.nn.Linear(2, 1, bias=False)
+    with torch.no_grad():
+        module.weight.copy_(torch.tensor([[1.75, -0.25]], dtype=torch.float32))
+    return module
+
+
+def _make_conv2d() -> torch.nn.Conv2d:
+    module = torch.nn.Conv2d(1, 1, kernel_size=2, bias=False)
+    with torch.no_grad():
+        module.weight.copy_(
+            torch.tensor([[[[1.75, -0.25], [0.5, -1.1]]]], dtype=torch.float32)
+        )
+    return module
+
+
+def test_quantize_maps_values_to_expected_levels() -> None:
+    x = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0], dtype=torch.float32)
+
+    actual = quantize(x, qmin=-2, qmax=2)
+    expected = torch.tensor([-3.0, -1.5, 0.0, 1.5, 3.0], dtype=torch.float32)
+
+    torch.testing.assert_close(actual, expected)
+
+
+def test_quantize_clamps_to_asymmetric_range() -> None:
+    x = torch.tensor([-2.0, -0.6, 0.2, 1.7], dtype=torch.float32)
+
+    actual = quantize(x, qmin=-1, qmax=2)
+    expected = torch.tensor([-1.0, -1.0, 0.0, 2.0], dtype=torch.float32)
+
+    torch.testing.assert_close(actual, expected)
+
+
+def test_quantize_preserves_straight_through_gradients() -> None:
+    x = torch.tensor([0.3, -0.6], dtype=torch.float32, requires_grad=True)
+    upstream = torch.tensor([2.0, -3.0], dtype=torch.float32)
+
+    loss = (quantize(x, qmin=-2, qmax=1) * upstream).sum()
+    loss.backward()
+
+    torch.testing.assert_close(x.grad, upstream)
+
+
+def test_quantize_zero_tensor_returns_zero_without_nan() -> None:
+    x = torch.zeros(4, dtype=torch.float32, requires_grad=True)
+
+    actual = quantize(x, qmin=-128, qmax=127)
+    actual.sum().backward()
+
+    torch.testing.assert_close(actual, torch.zeros_like(x))
+    torch.testing.assert_close(x.grad, torch.ones_like(x))
+    assert torch.isfinite(actual).all().item()
+
+
+def test_quantize_weight_applies_and_restores_on_linear() -> None:
+    module = _make_linear()
+    compressor = QuantizeWeight(min=-1, max=1)
+    x = torch.tensor([[2.0, 3.0]], dtype=torch.float32)
+    original_weight = module.weight.detach().clone()
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected_weight = quantize(original_weight, qmin=-1, qmax=1)
+    expected = F.linear(x, expected_weight)
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(module.weight.detach(), original_weight)
+
+
+def test_quantize_weight_applies_and_restores_on_conv2d() -> None:
+    module = _make_conv2d()
+    compressor = QuantizeWeight(min=-1, max=1)
+    x = torch.tensor([[[[1.0, 0.5], [-0.25, 0.75]]]], dtype=torch.float32)
+    original_weight = module.weight.detach().clone()
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected_weight = quantize(original_weight, qmin=-1, qmax=1)
+    expected = F.conv2d(x, expected_weight)
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(module.weight.detach(), original_weight)
+
+
+def test_quantize_activation_applies_to_linear_inputs() -> None:
+    module = _make_linear()
+    compressor = QuantizeActivation(min=-1, max=1)
+    x = torch.tensor([[1.75, -0.25]], dtype=torch.float32)
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected = F.linear(quantize(x, qmin=-1, qmax=1), module.weight.detach())
+
+    torch.testing.assert_close(actual, expected)
+
+    compressor.detach(module)
+    restored = module(x)
+    torch.testing.assert_close(restored, F.linear(x, module.weight.detach()))
+    assert not torch.equal(actual, restored)
+
+
+def test_quantize_activation_applies_to_conv2d_inputs() -> None:
+    module = _make_conv2d()
+    compressor = QuantizeActivation(min=-1, max=1)
+    x = torch.tensor([[[[1.75, -0.25], [0.5, -1.1]]]], dtype=torch.float32)
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected = F.conv2d(quantize(x, qmin=-1, qmax=1), module.weight.detach())
+
+    torch.testing.assert_close(actual, expected)
+
+    compressor.detach(module)
+    restored = module(x)
+    torch.testing.assert_close(restored, F.conv2d(x, module.weight.detach()))
+    assert not torch.equal(actual, restored)

rasnatune-0.0.1/tests/test_sparsification.py

@@ -0,0 +1,123 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from rasnatune.sparsification import (
+    SparseActivationUnstructured,
+    SparseWeightUnstructured,
+    sparsify,
+)
+
+
+def _make_linear() -> torch.nn.Linear:
+    module = torch.nn.Linear(4, 1, bias=False)
+    with torch.no_grad():
+        module.weight.copy_(torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.float32))
+    return module
+
+
+def _make_conv2d() -> torch.nn.Conv2d:
+    module = torch.nn.Conv2d(1, 1, kernel_size=2, bias=False)
+    with torch.no_grad():
+        module.weight.copy_(
+            torch.tensor([[[[1.0, -2.0], [3.0, -4.0]]]], dtype=torch.float32)
+        )
+    return module
+
+
+def test_sparsify_returns_input_for_zero_sparsity() -> None:
+    x = torch.tensor([1.0, -2.0, 3.0, -4.0], dtype=torch.float32)
+
+    actual = sparsify(x, sparsity=0.0)
+
+    torch.testing.assert_close(actual, x)
+
+
+def test_sparsify_zeroes_smallest_magnitudes() -> None:
+    x = torch.tensor([1.0, -2.0, 3.0, -4.0], dtype=torch.float32)
+
+    actual = sparsify(x, sparsity=0.5)
+    expected = torch.tensor([0.0, 0.0, 3.0, -4.0], dtype=torch.float32)
+
+    torch.testing.assert_close(actual, expected)
+
+
+def test_sparsify_zeros_all_values_for_full_sparsity() -> None:
+    x = torch.tensor([1.0, -2.0, 3.0, -4.0], dtype=torch.float32)
+
+    actual = sparsify(x, sparsity=1.0)
+
+    torch.testing.assert_close(actual, torch.zeros_like(x))
+
+
+def test_sparsify_rejects_invalid_sparsity() -> None:
+    with pytest.raises(ValueError):
+        sparsify(torch.ones(2, dtype=torch.float32), sparsity=1.5)
+
+
+def test_sparse_weight_applies_and_restores_on_linear() -> None:
+    module = _make_linear()
+    compressor = SparseWeightUnstructured(sparsity=0.5)
+    x = torch.tensor([[1.0, 2.0, 3.0, 4.0]], dtype=torch.float32)
+    original_weight = module.weight.detach().clone()
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected_weight = sparsify(original_weight, sparsity=0.5)
+    expected = F.linear(x, expected_weight)
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(module.weight.detach(), original_weight)
+
+
+def test_sparse_weight_applies_and_restores_on_conv2d() -> None:
+    module = _make_conv2d()
+    compressor = SparseWeightUnstructured(sparsity=0.5)
+    x = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], dtype=torch.float32)
+    original_weight = module.weight.detach().clone()
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected_weight = sparsify(original_weight, sparsity=0.5)
+    expected = F.conv2d(x, expected_weight)
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(module.weight.detach(), original_weight)
+
+
+def test_sparse_activation_applies_to_linear_inputs() -> None:
+    module = _make_linear()
+    compressor = SparseActivationUnstructured(sparsity=0.5)
+    x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.float32)
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected = F.linear(sparsify(x, sparsity=0.5), module.weight.detach())
+
+    torch.testing.assert_close(actual, expected)
+
+    compressor.detach(module)
+    restored = module(x)
+    torch.testing.assert_close(restored, F.linear(x, module.weight.detach()))
+    assert not torch.equal(actual, restored)
+
+
+def test_sparse_activation_applies_to_conv2d_inputs() -> None:
+    module = _make_conv2d()
+    compressor = SparseActivationUnstructured(sparsity=0.5)
+    x = torch.tensor([[[[1.0, -2.0], [3.0, -4.0]]]], dtype=torch.float32)
+
+    compressor.attach(module)
+    actual = module(x)
+
+    expected = F.conv2d(sparsify(x, sparsity=0.5), module.weight.detach())
+
+    torch.testing.assert_close(actual, expected)
+
+    compressor.detach(module)
+    restored = module(x)
+    torch.testing.assert_close(restored, F.conv2d(x, module.weight.detach()))
+    assert not torch.equal(actual, restored)