Type annotate poutines

[ordabayevy]

• escape_messenger
• indep_messenger
• infer_config_messenger
• lift_messenger
• mask_messenger
• uncondition_messenger

[ordabayevy]

Initially I tried to use class CondIndepStackFrame(NamedTuple), however, there is code in other places that tries to assign frame.full_size attribute which throws an error because NamedTuple is immutable(?).

[fritzo]

Hmm, immutability is a great property to reason about when reading & maintaining a codebase. I'd be a shame to lose that property even as we get type hints. Is there some way we could keep this class clean and immutable, even if that hack gets uglier? Even something like this:

@dataclass(frozen=True, slots=True)  # frozen + slots is close to namedtuple
class CondIndepStackFrame:
    name: str
    dim: Optional[int]
    size: int
    counter: int
    _full_size: Optional[int] = None  # hack: this is actually mutable

    @property
    def full_size(self) -> int:
        return self.size if self._full_size is None else self._full_size

    @full_size.setter
    def full_size(self, value: int) -> None:
        object.__setattr__(self, "_full_size", value)

    # ... __eq__ etc. that ignore .full_size

[ordabayevy]

I just tried this approach, however, python complained that you cannot use setattr with frame.full_size = size if frozen=True. Also slots is new in Python 3.10. I'll think more about other ways of keeping immutability, don't have any ideas at the moment.

[fritzo]

What if we kept CondIndepStackFrame as NamedTuple but moved FullSize to a global WeakKeyDictionary? There are only 4 code locations that use the hacky .full_size attribute. We could either change those to access the WeakKeyDictionary directly, or make a property to mediate access:

class CondIndepStackFrame(NamedTuple):
    name: str
    dim: Optional[int]
    size: int
    counter: int

    @property
    def full_size(self) -> Optional[int]:
        return COND_INDEP_FULL_SIZE.get(self)

    @full_size.getter
    def full_size(self, value: int) -> None:
        COND_INDEP_FULL_SIZE[self] = value

COND_INDEP_FULL_SIZE: WeakKeyDictionary[CondIndepStackFrame, int] = WeakKeyDictionary()

[fritzo]

And apologies for that full_size hack 😊

[ordabayevy]

I assume you meant @full_size.setter. Using NamedTuple and WeakKeyDictionary I get this error:

  File "/mnt/disks/dev/repos/pyro/pyro/poutine/subsample_messenger.py", line 154, in _process_message
    frame.full_size = self.size  # Used for param initialization.
  File "/mnt/disks/dev/repos/pyro/pyro/poutine/indep_messenger.py", line 28, in full_size
    COND_INDEP_FULL_SIZE[self] = value
  File "/home/yordabay/anaconda3/envs/pyro/lib/python3.8/weakref.py", line 396, in __setitem__
    self.data[ref(key, self._remove)] = value
TypeError: cannot create weak reference to 'CondIndepStackFrame' object

It seems like you cannot create weakrefs to tuples (https://stackoverflow.com/questions/58312618/is-there-a-way-to-support-weakrefs-with-collections-namedtuple)

I also tried it with a dataclass and using frozen=True but then it doesn't allow to use setter:

  File "/mnt/disks/dev/repos/pyro/pyro/poutine/subsample_messenger.py", line 154, in _process_message
    frame.full_size = self.size  # Used for param initialization.
  File "<string>", line 4, in __setattr__
dataclasses.FrozenInstanceError: cannot assign to field 'full_size'

[ordabayevy]

It seems like Python really doesn't want to allow setting attributes on immutable objects. It is also strange that an old namedtuple had no problems with that.

[fritzo]

Mind if I push a fix?

[ordabayevy]

Sure! That would be great. I don't have any solution on my end

[fritzo]

It's great to see such careful maintenance work, kudos!

I just have a couple comments about code clarity and immutability.

[fritzo]

Hmm, immutability is a great property to reason about when reading & maintaining a codebase. I'd be a shame to lose that property even as we get type hints. Is there some way we could keep this class clean and immutable, even if that hack gets uglier? Even something like this:

@dataclass(frozen=True, slots=True)  # frozen + slots is close to namedtuple
class CondIndepStackFrame:
    name: str
    dim: Optional[int]
    size: int
    counter: int
    _full_size: Optional[int] = None  # hack: this is actually mutable

    @property
    def full_size(self) -> int:
        return self.size if self._full_size is None else self._full_size

    @full_size.setter
    def full_size(self, value: int) -> None:
        object.__setattr__(self, "_full_size", value)

    # ... __eq__ etc. that ignore .full_size

[fritzo]

@ordabayevy your changes LGTM. Could you PTAL at my .full_size solution? Feel free to merge when tests pass.

[ordabayevy]

@fritzo LGTM. Thanks for the careful review!