google-deepmind/chex

post_init error in inherited dataclass

thomaspinder opened this issue · 1 comments

When inheriting one dataclass from another, Chex's dataclass does not allow a super() call to be made. This is something you can do in Python's base dataclass module.

A minimum working example is

from chex import dataclass as dataclass

@dataclass
class ChexBase:
    a : int 

    def __post_init__(self):
        self.b = self.a + 1

@dataclass
class ChexSub(ChexBase):
    a: int 

    def __post_init__(self):
        super().__post_init__()
        self.c = self.a + 2

temp = ChexSub(a = 1)
temp.b

Importing dataclass from dataclasses runs without error and returns 2, as expected.

Environment

  • Chex version 0.1.5
  • Ubuntu 20.04
  • Python 3.9

Hi @thomaspinder !

writing it this way

super(ChexSub, self).__post_init__()

solves the problem for me, maybe it'll help