Chex dataclass throws an exception in Python 3.9

[dionhaefner]

$ python --version
Python 3.9.1

In [1]: import chex

In [2]: @chex.dataclass
   ...: class Parameters:
   ...:   x: chex.ArrayDevice
   ...:   y: chex.ArrayDevice
   ...:
   ...: parameters = Parameters(
   ...:     x=jnp.ones((2, 2)),
   ...:     y=jnp.ones((1, 2)),
   ...: )
   ...:
   ...: # Dataclasses can be treated as JAX pytrees
   ...: jax.tree_map(lambda x: 2.0 * x, parameters)
   ...:
   ...: # and as mappings by dm-tree
   ...: tree.flatten(parameters)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-3-3461a2700932> in <module>
      1 @chex.dataclass
----> 2 class Parameters:
      3   x: chex.ArrayDevice
      4   y: chex.ArrayDevice
      5

~/.virtualenvs/science/lib/python3.9/site-packages/chex/_src/dataclass.py in dataclass(cls, init, repr, eq, order, unsafe_hash, frozen, mappable_dataclass, restricted_inheritance)
    104   if cls is None:
    105     return dcls
--> 106   return dcls(cls)
    107
    108

~/.virtualenvs/science/lib/python3.9/site-packages/chex/_src/dataclass.py in __call__(self, cls)
    147
    148     if self.mappable_dataclass:
--> 149       dcls = mappable_dataclass(dcls, self.restricted_inheritance)
    150
    151     def _from_tuple(args):

~/.virtualenvs/science/lib/python3.9/site-packages/chex/_src/dataclass.py in mappable_dataclass(cls, restricted_inheritance)
     81   if cls.__bases__ == (object,):
     82     # `collections.Mapping` is incompatible with `object`
---> 83     cls.__bases__ = (collections.Mapping,)
     84   else:
     85     cls.__bases__ += (collections.Mapping,)

TypeError: __bases__ assignment: 'Mapping' deallocator differs from 'object'

[hbq1]

Thanks a lot for filing this bug! It should be fixed now.