Skip to main content

Features

Pyre has custom support for Python idioms that would otherwise not be supported by the usual type annotations.

Registering attributes using PyTorch's register_buffer#

PyTorch allows subclasses of nn.Module to register a buffer in an object using self.register_buffer("foo", initial_value). Pyre supports this pattern when used within the constructor. It simply treats the buffer as a Tensor attribute of the class:

import torchimport torch.nn as nn
class Foo(nn.Module):    def __init__(self) -> None:        super(Foo, self).__init__()        self.register_buffer("foo", torch.zeros(10, 20))        self.register_buffer("foo_persistent", torch.zeros(10, 20), persistent=False)
    def bar(self) -> None:        reveal_type(self.foo) # => torch.Tensor        reveal_type(self.foo_persistent) # => torch.Tensor
def baz() -> None:    y = Foo().foo    reveal_type(y) # => torch.Tensor

Note that Pyre will not recognize buffers registered in methods other than the constructor (just like it doesn't recognize attributes defined in methods other than the constructor).

It will also not recognize buffers that are initialized with None since it cannot infer the exact type of the buffer. In such a case, you can tell Pyre about the attribute's type by explicitly defining it in the class:

import torchimport torch.nn as nnfrom typing import Optional
class Foo(nn.Module):    my_none_buffer: Optional[torch.Tensor]
    def __init__(self) -> None:        super(Foo, self).__init__()        self.register_buffer("my_none_buffer", None)
    def bar(self) -> None:        reveal_type(self.my_none_buffer) # => Optional[torch.Tensor]