tencheck 0.0.2__tar.gz

tencheck-0.0.2/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2024-present Justin Yan
+
+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.

tencheck-0.0.2/PKG-INFO

@@ -0,0 +1,49 @@
+Metadata-Version: 2.1
+Name: tencheck
+Version: 0.0.2
+Summary: A library for pytorch layer testing.
+Author-email: Justin Yan <justin@iomorphic.com>
+Project-URL: Homepage, https://github.com/justin-yan/tencheck
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<2.4,>2
+Requires-Dist: jaxtyping
+Requires-Dist: numpy
+
+# tencheck
+
+`tencheck` provides a simple set of utilities for analyzing and validating the properties of layers of deep neural nets.
+
+It's typically quite difficult to validate that a layer "behaves properly" (oftentimes, the final barometer is simply how well a model performs with the layer included), and many brittle unit tests have been written involving randomly instantiated tensors and `is_close` checks.  We believe a good "first line of defense" for neural nets is to create a suite of properties that can be asserted about a layer, while requiring minimal effort per layer in order to do so.
+
+We think there are two aspects of property-based testing that are quite useful to take inspiration from:
+
+- Automatically generating inputs (and generating inputs of variable sizes and values to elucidate properties of interest).
+- Evaluating properties based on the maintenance of invariants instead of attempting to exactly match values (which is particularly difficult to interpret in deep neural nets).
+
+However, an important difference is that the properties of interest are generally fairly generic and often shared between layers, while the input generation strategies are pretty similar (they're all tensors).  So the focus of `tencheck` is to provide:
+
+- An (attempted) universal input generation harness.
+- A variety of interesting properties.
+- Three modalities: assertion, analysis, and profiling.
+
+The following requirements need to be met for `tencheck` to work:
+
+- Your layers are implemented in torch.
+- The `.forward()` method is annotated with [jaxtyping](https://github.com/patrick-kidger/jaxtyping).
+
+
+## Backlog
+
+- For profiling, use a grid of input sizes to generate performance curves.
+- Pick a flop counter and use for profiling
+- Tensor container types include more options like dataclasses.
+- Auto-generate simple hyperparameters for layer instantiation.
+- Refine dtype mapping and coherence.

tencheck-0.0.2/README.md

@@ -0,0 +1,30 @@
+# tencheck
+
+`tencheck` provides a simple set of utilities for analyzing and validating the properties of layers of deep neural nets.
+
+It's typically quite difficult to validate that a layer "behaves properly" (oftentimes, the final barometer is simply how well a model performs with the layer included), and many brittle unit tests have been written involving randomly instantiated tensors and `is_close` checks.  We believe a good "first line of defense" for neural nets is to create a suite of properties that can be asserted about a layer, while requiring minimal effort per layer in order to do so.
+
+We think there are two aspects of property-based testing that are quite useful to take inspiration from:
+
+- Automatically generating inputs (and generating inputs of variable sizes and values to elucidate properties of interest).
+- Evaluating properties based on the maintenance of invariants instead of attempting to exactly match values (which is particularly difficult to interpret in deep neural nets).
+
+However, an important difference is that the properties of interest are generally fairly generic and often shared between layers, while the input generation strategies are pretty similar (they're all tensors).  So the focus of `tencheck` is to provide:
+
+- An (attempted) universal input generation harness.
+- A variety of interesting properties.
+- Three modalities: assertion, analysis, and profiling.
+
+The following requirements need to be met for `tencheck` to work:
+
+- Your layers are implemented in torch.
+- The `.forward()` method is annotated with [jaxtyping](https://github.com/patrick-kidger/jaxtyping).
+
+
+## Backlog
+
+- For profiling, use a grid of input sizes to generate performance curves.
+- Pick a flop counter and use for profiling
+- Tensor container types include more options like dataclasses.
+- Auto-generate simple hyperparameters for layer instantiation.
+- Refine dtype mapping and coherence.

tencheck-0.0.2/pyproject.toml

@@ -0,0 +1,90 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tencheck"
+version = "0.0.2"
+authors = [
+    { name="Justin Yan", email="justin@iomorphic.com" }
+]
+description = "A library for pytorch layer testing."
+readme = "README.md"
+requires-python = ">=3.11"
+classifiers = [
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+]
+dependencies = [
+    ######
+    ### Custom Dependencies Section Begin
+    ######
+    "torch>2,<2.4",
+    "jaxtyping",
+    "numpy",
+    ######
+    ### Custom Dependencies Section End
+    ######
+]
+
+[project.urls]
+"Homepage" = "https://github.com/justin-yan/tencheck"
+
+[tool.setuptools]
+zip-safe = false
+include-package-data = true
+
+[tool.setuptools.package-data]
+"tencheck" = ["py.typed"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+#######
+### Miscellaneous Tool Configuration
+#######
+[tool.ruff]
+line-length = 150
+target-version = "py311"
+
+[tool.ruff.format]
+quote-style = "double"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I"]
+ignore = ["F722"]
+
+[tool.ruff.lint.isort]
+known-first-party = ["tencheck"]
+
+[tool.pytest.ini_options]
+addopts = "-ra -q --doctest-modules --jaxtyping-packages=tencheck.examples,beartype.beartype"
+
+[tool.mypy]
+mypy_path = "src"
+disallow_untyped_defs = true
+disallow_any_unimported = true
+allow_redefinition = false
+ignore_errors = false
+implicit_reexport = false
+local_partial_types = true
+no_implicit_optional = true
+strict_equality = true
+strict_optional = true
+warn_no_return = true
+warn_redundant_casts = true
+warn_unreachable = true
+warn_unused_configs = true
+warn_unused_ignores = true
+
+######
+### Custom Directives Section Begin
+######
+
+######
+### Custom Directives Section End
+######

tencheck-0.0.2/setup.cfg

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

tencheck-0.0.2/src/tencheck/checks.py

@@ -0,0 +1,15 @@
+from torch import nn as nn
+
+
+def unused_params_check(layer: nn.Module) -> None:
+    """
+    This check is run after a backward pass is completed.
+
+    If any unused parameters are found, an exception is thrown with the named parameters.
+    """
+    unused_parameters = []
+    for name, param in layer.named_parameters():
+        if param.grad is None:
+            unused_parameters.append(name)
+
+    assert len(unused_parameters) == 0, f"Unused parameters: {unused_parameters} detected."

tencheck-0.0.2/src/tencheck/examples.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+
+
+class SimpleLinReluModule(nn.Module):
+    def __init__(self, out_features: int) -> None:
+        super(SimpleLinReluModule, self).__init__()
+        self.linear = nn.Linear(32, out_features)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: Float[torch.Tensor, "B 32"]) -> Float[torch.Tensor, "B O"]:
+        x = self.linear(x)
+        x = self.relu(x)
+        return x
+
+
+class CasedLinReluModule(nn.Module):
+    _tencheck_cases = [{"out_features": 10}, {"out_features": 20}]
+
+    def __init__(self, out_features: int) -> None:
+        super(CasedLinReluModule, self).__init__()
+        self.linear = nn.Linear(32, out_features)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: Float[torch.Tensor, "B 32"]) -> Float[torch.Tensor, "B O"]:
+        x = self.linear(x)
+        x = self.relu(x)
+        return x
+
+
+class UnusedParamsModule(nn.Module):
+    def __init__(self, out_features: int) -> None:
+        super(UnusedParamsModule, self).__init__()
+        self.linear = nn.Linear(32, out_features)
+        self.unused_linear = nn.Linear(32, out_features)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: Float[torch.Tensor, "B 32"]) -> Float[torch.Tensor, "B O"]:
+        x = self.linear(x)
+        x = self.relu(x)
+        return x
+
+
+class BrokenModule(nn.Module):
+    def __init__(self) -> None:
+        super(BrokenModule, self).__init__()
+        self.linear = nn.Linear(32, 4)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: Float[torch.Tensor, "B 32"]) -> Float[torch.Tensor, "B 4"]:
+        x = self.linear(x)
+        raise Exception("Module is broken")
+        x = self.relu(x)  # type: ignore[unreachable]
+        return x
+
+
+class MistypedModule(nn.Module):
+    def __init__(self, out_features: int) -> None:
+        super(MistypedModule, self).__init__()
+        self.linear = nn.Linear(32, out_features)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: Float[torch.Tensor, "B 32"]) -> Float[torch.Tensor, "B 1025"]:
+        x = self.linear(x)
+        x = self.relu(x)
+        return x

tencheck-0.0.2/src/tencheck/harness.py

@@ -0,0 +1,68 @@
+import time
+from typing import Optional, Type
+
+import torch
+from torch import nn as nn
+
+from tencheck.checks import unused_params_check
+from tencheck.input import input_gen
+from tencheck.loss import trivial_loss
+from tencheck.ttypes import CaseDefined, LayerStats
+
+
+def check_layers(layers: list[nn.Module | Type[CaseDefined]], seed: Optional[int] = None) -> None:
+    """
+    This method receives a *concrete* list of layer objects, and asserts the relevant properties.
+    """
+    for layer in layers:
+        if isinstance(layer, CaseDefined):  # This works even though layer is a class, not an obj, due to @runtime_checkable
+            for case in layer._tencheck_cases:
+                layer_obj = layer(**case)
+                _single_layer_assert_all(layer_obj, seed)
+        else:
+            _single_layer_assert_all(layer, seed)  # type: ignore[arg-type]
+
+
+def _single_layer_assert_all(layer: nn.Module, seed: Optional[int] = None) -> None:
+    in_tens = input_gen(layer, seed)
+    layer.zero_grad(set_to_none=True)
+
+    # throws TypeCheckError for shapecheck
+    # throws Exception for generic issues
+    out = layer.forward(**in_tens)
+    loss = trivial_loss(out)
+    loss.backward()
+
+    unused_params_check(layer)
+
+
+def profile_layer(layer: nn.Module) -> LayerStats:
+    """
+    This runs a basic profiling setup for a single layer.
+    """
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        raise Exception("GPU is not available")
+
+    layer = layer.to(device)
+    layer.zero_grad(set_to_none=True)
+    in_tens = input_gen(layer, device=device)
+
+    torch.cuda.reset_peak_memory_stats()
+
+    start = time.perf_counter()
+    out = layer.forward(**in_tens)
+    loss = trivial_loss(out)
+    loss.backward()
+    elapsed = time.perf_counter() - start
+    peak_memory_gbs = torch.cuda.max_memory_allocated() / 1024**3
+    gigaflops = 0.0
+
+    return LayerStats(elapsed, peak_memory_gbs, gigaflops)
+
+
+if __name__ == "__main__":
+    from tencheck.examples import SimpleLinReluModule
+
+    print(profile_layer(SimpleLinReluModule(5)))

tencheck-0.0.2/src/tencheck/input.py

@@ -0,0 +1,84 @@
+import inspect
+import logging
+import random
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from jaxtyping._array_types import _FixedDim, _NamedDim, _NamedVariadicDim, _SymbolicDim
+from torch import Tensor
+from torch.testing import make_tensor
+
+dtype_mapping = {
+    "bool": torch.bool,
+    "bool_": torch.bool,
+    "uint4": torch.int32,
+    "uint8": torch.int32,
+    "uint16": torch.int32,
+    "uint32": torch.int32,
+    "uint64": torch.int32,
+    "int4": torch.int32,
+    "int8": torch.int32,
+    "int16": torch.int32,
+    "int32": torch.int32,
+    "int64": torch.int32,
+    "bfloat16": torch.float32,
+    "float16": torch.float32,
+    "float32": torch.float32,
+    "float64": torch.float32,
+    "complex64": torch.float32,
+    "complex128": torch.float32,
+}
+
+
+@dataclass
+class TensorSpec:
+    shape: tuple[int, ...]
+    dtype: torch.dtype
+
+
+def input_gen(layer: nn.Module, seed: Optional[int] = None, device: str | torch.device = "cpu") -> dict[str, Tensor]:
+    """
+    For a given layer that is type annotated with jaxtyping, produce a map of mock tensors that can be used like so:
+
+    in_tens = input_gen(layer)
+    layer.forward(**in_tens)
+    """
+    if seed:
+        torch.manual_seed(seed)
+        random.seed(seed)
+
+    signature = inspect.signature(layer.forward)
+    tensor_specs: dict[str, TensorSpec] = {}  # parameter to shape
+    # Across all of the parameters, we'll have a mix of `_NamedDim`, `_FixedDim`, `_NamedVariadicDim`, or `_SymbolicDim`
+    # We want to pick concrete dimensions for everything that isn't fixed,
+    #   and then we want to generate tensors for everything.
+    dimension_name_map: dict[str, int] = {}
+    for name, param_obj in signature.parameters.items():
+        shape: list[int] = []
+        for dim in param_obj.annotation.dims:
+            match dim:
+                case _NamedDim(nm, _, _):
+                    sz = dimension_name_map.setdefault(nm, random.randint(4, 64))
+                    shape.append(sz)
+                case _FixedDim(sz, _):
+                    shape.append(sz)
+                case _NamedVariadicDim() | _SymbolicDim() | _:
+                    raise NotImplementedError("Don't yet handle these dimension cases")
+        dt = dtype_mapping[param_obj.annotation.dtypes[0]]
+        tensor_specs[name] = TensorSpec(tuple(shape), dt)
+
+    output = {}
+    for name, spec in tensor_specs.items():
+        match spec.dtype:
+            case torch.float32:
+                mock_ten = make_tensor(spec.shape, dtype=spec.dtype, device=device, low=-1, high=1)
+            case torch.int32:
+                mock_ten = make_tensor(spec.shape, dtype=spec.dtype, device=device, low=-10, high=10)
+            case _:
+                logging.error(spec)
+                raise NotImplementedError("Don't yet handle these dtypes")
+        output[name] = mock_ten
+
+    return output

tencheck-0.0.2/src/tencheck/loss.py

@@ -0,0 +1,34 @@
+import torch
+from torch import Tensor
+
+TensorContainerTypes = Tensor | list | set | tuple | dict
+
+
+def flatten_tensors(tct: TensorContainerTypes) -> list[Tensor]:
+    """
+    This method will recursively traverse any potential container structures and extract any tensors as a flat list.
+    """
+    tensor_list: list[Tensor] = []
+    match tct:
+        case Tensor():
+            tensor_list.append(tct)
+        case list():
+            for e in tct:
+                match e:
+                    case Tensor():
+                        tensor_list.append(e)
+                    case _:
+                        tensor_list.extend(flatten_tensors(e))
+        case set() | tuple():
+            tensor_list.extend(flatten_tensors(list(tct)))
+        case dict():
+            tensor_list.extend(flatten_tensors([v for k, v in tct.items()]))
+        case _:
+            raise NotImplementedError("Unexpected input format")
+
+    return tensor_list
+
+
+def trivial_loss(tct: TensorContainerTypes) -> Tensor:
+    partial_sums = [t.sum() for t in flatten_tensors(tct)]
+    return torch.stack(partial_sums).sum()

tencheck-0.0.2/src/tencheck/ttypes.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class CaseDefined(Protocol):
+    @property
+    def _tencheck_cases(self) -> List[Dict[str, Any]]:
+        raise NotImplementedError()
+
+
+@dataclass
+class LayerStats:
+    total_time: float
+    peak_mem_gigs: float
+    giga_flop_count: float

tencheck-0.0.2/src/tencheck.egg-info/PKG-INFO

@@ -0,0 +1,49 @@
+Metadata-Version: 2.1
+Name: tencheck
+Version: 0.0.2
+Summary: A library for pytorch layer testing.
+Author-email: Justin Yan <justin@iomorphic.com>
+Project-URL: Homepage, https://github.com/justin-yan/tencheck
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<2.4,>2
+Requires-Dist: jaxtyping
+Requires-Dist: numpy
+
+# tencheck
+
+`tencheck` provides a simple set of utilities for analyzing and validating the properties of layers of deep neural nets.
+
+It's typically quite difficult to validate that a layer "behaves properly" (oftentimes, the final barometer is simply how well a model performs with the layer included), and many brittle unit tests have been written involving randomly instantiated tensors and `is_close` checks.  We believe a good "first line of defense" for neural nets is to create a suite of properties that can be asserted about a layer, while requiring minimal effort per layer in order to do so.
+
+We think there are two aspects of property-based testing that are quite useful to take inspiration from:
+
+- Automatically generating inputs (and generating inputs of variable sizes and values to elucidate properties of interest).
+- Evaluating properties based on the maintenance of invariants instead of attempting to exactly match values (which is particularly difficult to interpret in deep neural nets).
+
+However, an important difference is that the properties of interest are generally fairly generic and often shared between layers, while the input generation strategies are pretty similar (they're all tensors).  So the focus of `tencheck` is to provide:
+
+- An (attempted) universal input generation harness.
+- A variety of interesting properties.
+- Three modalities: assertion, analysis, and profiling.
+
+The following requirements need to be met for `tencheck` to work:
+
+- Your layers are implemented in torch.
+- The `.forward()` method is annotated with [jaxtyping](https://github.com/patrick-kidger/jaxtyping).
+
+
+## Backlog
+
+- For profiling, use a grid of input sizes to generate performance curves.
+- Pick a flop counter and use for profiling
+- Tensor container types include more options like dataclasses.
+- Auto-generate simple hyperparameters for layer instantiation.
+- Refine dtype mapping and coherence.

tencheck-0.0.2/src/tencheck.egg-info/SOURCES.txt

@@ -0,0 +1,17 @@
+LICENSE
+README.md
+pyproject.toml
+src/tencheck/__init__.py
+src/tencheck/checks.py
+src/tencheck/examples.py
+src/tencheck/harness.py
+src/tencheck/input.py
+src/tencheck/loss.py
+src/tencheck/py.typed
+src/tencheck/ttypes.py
+src/tencheck.egg-info/PKG-INFO
+src/tencheck.egg-info/SOURCES.txt
+src/tencheck.egg-info/dependency_links.txt
+src/tencheck.egg-info/not-zip-safe
+src/tencheck.egg-info/requires.txt
+src/tencheck.egg-info/top_level.txt

tencheck-0.0.2/src/tencheck.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

tencheck-0.0.2/src/tencheck.egg-info/not-zip-safe

@@ -0,0 +1 @@
+

tencheck-0.0.2/src/tencheck.egg-info/requires.txt

@@ -0,0 +1,3 @@
+torch<2.4,>2
+jaxtyping
+numpy

tencheck-0.0.2/src/tencheck.egg-info/top_level.txt

@@ -0,0 +1 @@
+tencheck