DOCSEARCHbyAlgolia
Submit Docs

torchtyping

repository·master·Indexed 23 days ago

https://github.com/patrick-kidger/torchtyping

A library providing type annotations for PyTorch tensors to specify shape, dtype, and layout constraints. It supports runtime enforcement of these constraints via typeguard and provides the TensorType API for defining fixed dimensions, named dimensions, and custom tensor details.

Tokens
1.8K
Snippets
5
Records
10
Agent score
31%

What's inside torchtyping

  1. How to create custom TensorDetail extensions

    master

    You can extend torchtyping to check arbitrary properties of a tensor by subclassing torchtyping.TensorDetail. A custom detail must implement three methods:

    1. check(self, tensor: Tensor) -> bool: Returns True if the tensor satisfies the condition.
    2. __repr__(self) -> str: Returns a string describing the detail (used in error messages).
    3. tensor_repr(cls, tensor: Tensor) -> str: Returns a string describing the actual property found on the tensor (used in error messages).

    Pass instances of your custom detail to TensorType to enforce the check.

    from torch import rand, Tensor
from torchtyping import TensorDetail, TensorType
from typeguard import typechecked

class FooDetail(TensorDetail):
    def __init__(self, foo):
        super().__init__()
        self.foo = foo
        
    def check(self, tensor: Tensor) -> bool:
        return hasattr(tensor, "foo") and tensor.foo == self.foo

    def __repr__(self) -> str:
        return f"FooDetail({self.foo})"

    @classmethod
    def tensor_repr(cls, tensor: Tensor) -> str:
        if hasattr(tensor, "foo"):
            return f"FooDetail({tensor.foo})"
       	else:
            return ""

@typechecked
def foo_checker(tensor: TensorType[float, FooDetail("good-foo")]):
    pass

  2. Enable runtime type checking with typeguard

    master

    To enable runtime type checking, you must call torchtyping.patch_typeguard() before defining the functions you wish to check. Additionally, you must enable typeguard using one of the following methods:

    1. Decorate the function with @typeguard.typechecked.
    2. Use typeguard.importhook.install_import_hook().
    3. Use pytest command line flags (as described in the main README).

    Important: Ensure the functions are defined after the call to patch_typeguard().

  3. Configure mypy for torchtyping

    master

    Because torchtyping uses functionality currently beyond what mypy can represent (PEP 646), you must tell mypy to ignore the import statements to avoid errors:

    from torchtyping import TensorType  # type: ignore

    Mypy Crash Workaround: mypy may crash on files using the str: int or str: ... notation (e.g., TensorType["batch": 10]). To work around this, create a .pyi stub file for that specific filename in the same directory.

  7. Use TensorType for shape and dtype annotations

    master

    Use TensorType to annotate PyTorch tensors with shape, dtype, layout, and other details. This provides clear documentation and, if typeguard is installed, runtime enforcement of these constraints.

    Example of annotating a batch outer product where dimensions are named and checked for consistency:

    from torch import Tensor
from torchtyping import TensorType

def batch_outer_product(x: TensorType["batch", "x_channels"],
                        y: TensorType["batch", "y_channels"]
                        ) -> TensorType["batch", "x_channels", "y_channels"]:
    return x.unsqueeze(-1) * y.unsqueeze(-2)

  8. Use patch_typeguard()

    master

    The torchtyping.patch_typeguard() function integrates torchtyping with typeguard to perform runtime checks. It is safe to call multiple times.

    • If using @typeguard.typechecked: Call it before using the decorator.
    • If using typeguard.importhook.install_import_hook: Call it any time before defining checked functions.
    • If not using typeguard: This function can be omitted.

  9. Common TensorType syntax patterns

    master

    Use these patterns to define various tensor constraints:

    • Fixed dimensions: TensorType[3, 4] (shape (3, 4))
    • Any size dimension: TensorType[2, -1, -1] (shape (2, any, any))
    • Named dimensions: TensorType["batch", 5] (shape (batch, 5))
    • Ellipsis (arbitrary batch dims): TensorType[..., 2, 3] (shape (..., 2, 3))
    • Named batch dimensions: TensorType["batch": ..., "channels_x"] (checks that the ... dimensions match across arguments)
    • Scalar (zero-dim): TensorType[()]
    • Dtype only: TensorType[float] (matches torch.get_default_dtype())
    • Shape and Dtype: TensorType[3, 4, float]
    • Layout: TensorType[torch.sparse_coo]
    • Named Tensors: TensorType["a": 3, "b", is_named]

  10. Reference the TensorType API

    master

    The core API is torchtyping.TensorType[shape, dtype, layout, details]. All arguments are optional.

    Shape Argument

    • int: Exact dimension size. -1 allows any size.
    • str: Binds the dimension size to a name for consistency checking across tensors.
    • ...: Represents an arbitrary number of dimensions.
    • str: int: A slice combining name and exact size.
    • str: str: Binds a dimension to two names (useful for documentation).
    • str: ...: Binds multiple dimensions to a specific name.
    • None: Indicates a dimension that must NOT have a name (per PyTorch named tensors).
    • None: int / None: str: Combines None with size or name constraints.
    • typing.Any: Any size allowed (equivalent to -1).
    • tuple: A combination of the above, e.g., TensorType["batch": ..., "length": 10].

    Dtype Argument

    • torch.float32, torch.float64, etc.
    • int, bool, float (converts to corresponding PyTorch types; float uses torch.get_default_dtype()).

    Layout Argument

    • torch.strided (dense)
    • torch.sparse_coo (sparse)

    Details Argument

    Allows passing additional flags for customization. Built-in flags include:

    • torchtyping.is_named: Checks if tensor dimensions have names.
    • torchtyping.is_float: Checks if the tensor is an arbitrary floating point type.